|
1 #!/usr/bin/python2.4 |
|
2 # |
|
3 # Copyright 2008 Google Inc. |
|
4 # |
|
5 # Licensed under the Apache License, Version 2.0 (the "License"); |
|
6 # you may not use this file except in compliance with the License. |
|
7 # You may obtain a copy of the License at |
|
8 # |
|
9 # http://www.apache.org/licenses/LICENSE-2.0 |
|
10 # |
|
11 # Unless required by applicable law or agreed to in writing, software |
|
12 # distributed under the License is distributed on an "AS IS" BASIS, |
|
13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|
14 # See the License for the specific language governing permissions and |
|
15 # limitations under the License. |
|
16 |
|
17 import inspect |
|
18 |
|
19 class StubOutForTesting: |
|
20 """Sample Usage: |
|
21 You want os.path.exists() to always return true during testing. |
|
22 |
|
23 stubs = StubOutForTesting() |
|
24 stubs.Set(os.path, 'exists', lambda x: 1) |
|
25 ... |
|
26 stubs.UnsetAll() |
|
27 |
|
28 The above changes os.path.exists into a lambda that returns 1. Once |
|
29 the ... part of the code finishes, the UnsetAll() looks up the old value |
|
30 of os.path.exists and restores it. |
|
31 |
|
32 """ |
|
33 def __init__(self): |
|
34 self.cache = [] |
|
35 self.stubs = [] |
|
36 |
|
37 def __del__(self): |
|
38 self.SmartUnsetAll() |
|
39 self.UnsetAll() |
|
40 |
|
41 def SmartSet(self, obj, attr_name, new_attr): |
|
42 """Replace obj.attr_name with new_attr. This method is smart and works |
|
43 at the module, class, and instance level while preserving proper |
|
44 inheritance. It will not stub out C types however unless that has been |
|
45 explicitly allowed by the type. |
|
46 |
|
47 This method supports the case where attr_name is a staticmethod or a |
|
48 classmethod of obj. |
|
49 |
|
50 Notes: |
|
51 - If obj is an instance, then it is its class that will actually be |
|
52 stubbed. Note that the method Set() does not do that: if obj is |
|
53 an instance, it (and not its class) will be stubbed. |
|
54 - The stubbing is using the builtin getattr and setattr. So, the __get__ |
|
55 and __set__ will be called when stubbing (TODO: A better idea would |
|
56 probably be to manipulate obj.__dict__ instead of getattr() and |
|
57 setattr()). |
|
58 |
|
59 Raises AttributeError if the attribute cannot be found. |
|
60 """ |
|
61 if (inspect.ismodule(obj) or |
|
62 (not inspect.isclass(obj) and obj.__dict__.has_key(attr_name))): |
|
63 orig_obj = obj |
|
64 orig_attr = getattr(obj, attr_name) |
|
65 |
|
66 else: |
|
67 if not inspect.isclass(obj): |
|
68 mro = list(inspect.getmro(obj.__class__)) |
|
69 else: |
|
70 mro = list(inspect.getmro(obj)) |
|
71 |
|
72 mro.reverse() |
|
73 |
|
74 orig_attr = None |
|
75 |
|
76 for cls in mro: |
|
77 try: |
|
78 orig_obj = cls |
|
79 orig_attr = getattr(obj, attr_name) |
|
80 except AttributeError: |
|
81 continue |
|
82 |
|
83 if orig_attr is None: |
|
84 raise AttributeError("Attribute not found.") |
|
85 |
|
86 # Calling getattr() on a staticmethod transforms it to a 'normal' function. |
|
87 # We need to ensure that we put it back as a staticmethod. |
|
88 old_attribute = obj.__dict__.get(attr_name) |
|
89 if old_attribute is not None and isinstance(old_attribute, staticmethod): |
|
90 orig_attr = staticmethod(orig_attr) |
|
91 |
|
92 self.stubs.append((orig_obj, attr_name, orig_attr)) |
|
93 setattr(orig_obj, attr_name, new_attr) |
|
94 |
|
95 def SmartUnsetAll(self): |
|
96 """Reverses all the SmartSet() calls, restoring things to their original |
|
97 definition. Its okay to call SmartUnsetAll() repeatedly, as later calls |
|
98 have no effect if no SmartSet() calls have been made. |
|
99 |
|
100 """ |
|
101 self.stubs.reverse() |
|
102 |
|
103 for args in self.stubs: |
|
104 setattr(*args) |
|
105 |
|
106 self.stubs = [] |
|
107 |
|
108 def Set(self, parent, child_name, new_child): |
|
109 """Replace child_name's old definition with new_child, in the context |
|
110 of the given parent. The parent could be a module when the child is a |
|
111 function at module scope. Or the parent could be a class when a class' |
|
112 method is being replaced. The named child is set to new_child, while |
|
113 the prior definition is saved away for later, when UnsetAll() is called. |
|
114 |
|
115 This method supports the case where child_name is a staticmethod or a |
|
116 classmethod of parent. |
|
117 """ |
|
118 old_child = getattr(parent, child_name) |
|
119 |
|
120 old_attribute = parent.__dict__.get(child_name) |
|
121 if old_attribute is not None and isinstance(old_attribute, staticmethod): |
|
122 old_child = staticmethod(old_child) |
|
123 |
|
124 self.cache.append((parent, old_child, child_name)) |
|
125 setattr(parent, child_name, new_child) |
|
126 |
|
127 def UnsetAll(self): |
|
128 """Reverses all the Set() calls, restoring things to their original |
|
129 definition. Its okay to call UnsetAll() repeatedly, as later calls have |
|
130 no effect if no Set() calls have been made. |
|
131 |
|
132 """ |
|
133 # Undo calls to Set() in reverse order, in case Set() was called on the |
|
134 # same arguments repeatedly (want the original call to be last one undone) |
|
135 self.cache.reverse() |
|
136 |
|
137 for (parent, old_child, child_name) in self.cache: |
|
138 setattr(parent, child_name, old_child) |
|
139 self.cache = [] |