OLD | NEW |
| (Empty) |
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: | |
122 if isinstance(old_attribute, staticmethod): | |
123 old_child = staticmethod(old_child) | |
124 elif isinstance(old_attribute, classmethod): | |
125 old_child = classmethod(old_child.im_func) | |
126 | |
127 self.cache.append((parent, old_child, child_name)) | |
128 setattr(parent, child_name, new_child) | |
129 | |
130 def UnsetAll(self): | |
131 """Reverses all the Set() calls, restoring things to their original | |
132 definition. Its okay to call UnsetAll() repeatedly, as later calls have | |
133 no effect if no Set() calls have been made. | |
134 | |
135 """ | |
136 # Undo calls to Set() in reverse order, in case Set() was called on the | |
137 # same arguments repeatedly (want the original call to be last one undone) | |
138 self.cache.reverse() | |
139 | |
140 for (parent, old_child, child_name) in self.cache: | |
141 setattr(parent, child_name, old_child) | |
142 self.cache = [] | |
OLD | NEW |