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 |