• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2008 Google Inc.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#      http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""
16This is a fork of the pymox library intended to work with Python 3.
17The file was modified by quermit@gmail.com and dawid.fatyga@gmail.com
18
19Previously, pyfakefs used just this file from the mox3 library.
20However, mox3 will soon be decommissioned, yet standard mock cannot
21be used because of the problem described in pyfakefs #182 and
22mock issue 250 (https://github.com/testing-cabal/mock/issues/250).
23Therefore just this file was forked from mox3 and incorporated
24into pyfakefs.
25"""
26
27import inspect
28
29
30class StubOutForTesting:
31    """Sample Usage:
32
33    You want os.path.exists() to always return true during testing.
34
35    stubs = StubOutForTesting()
36    stubs.Set(os.path, 'exists', lambda x: 1)
37        ...
38    stubs.UnsetAll()
39
40    The above changes os.path.exists into a lambda that returns 1.    Once
41    the ... part of the code finishes, the UnsetAll() looks up the old value
42    of os.path.exists and restores it.
43
44    """
45
46    def __init__(self):
47        self.cache = []
48        self.stubs = []
49
50    def __del__(self):
51        self.smart_unset_all()
52        self.unset_all()
53
54    def smart_set(self, obj, attr_name, new_attr):
55        """Replace obj.attr_name with new_attr.
56
57        This method is smart and works at the module, class, and instance level
58        while preserving proper inheritance. It will not stub out C types
59        however unless that has been explicitly allowed by the type.
60
61        This method supports the case where attr_name is a staticmethod or a
62        classmethod of obj.
63
64        If obj is an instance, then it is its class that will actually be
65        stubbed. Note that the method Set() does not do that: if obj is an
66        instance, it (and not its class) will be stubbed.
67
68        Raises AttributeError if the attribute cannot be found.
69        """
70        if inspect.ismodule(obj) or (
71            not inspect.isclass(obj) and attr_name in obj.__dict__
72        ):
73            orig_obj = obj
74            if attr_name in obj.__dict__:
75                orig_attr = obj.__dict__[attr_name]
76            else:
77                orig_attr = None
78
79        else:
80            if not inspect.isclass(obj):
81                mro = list(inspect.getmro(obj.__class__))
82            else:
83                mro = list(inspect.getmro(obj))
84
85            mro.reverse()
86
87            orig_attr = None
88
89            for cls in mro:
90                try:
91                    orig_obj = cls
92                    orig_attr = obj.__dict__[attr_name]
93                except KeyError:
94                    continue
95
96        if orig_attr is None:
97            raise AttributeError("Attribute not found.")
98
99        self.stubs.append((orig_obj, attr_name, orig_attr))
100        orig_obj.__dict__[attr_name] = new_attr
101
102    def smart_unset_all(self):
103        """Reverses all the SmartSet() calls.
104
105        Restores things to their original definition. Its okay to call
106        SmartUnsetAll() repeatedly, as later calls have no effect if no
107        SmartSet() calls have been made.
108        """
109        self.stubs.reverse()
110
111        for obj, attr_name, old_attr in self.stubs:
112            obj.__dict__[attr_name] = old_attr
113
114        self.stubs = []
115
116    def set(self, parent, child_name, new_child):
117        """Replace child_name's old definition with new_child.
118
119        Replace definition in the context of the given parent. The parent could
120        be a module when the child is a function at module scope. Or the parent
121        could be a class when a class' method is being replaced. The named
122        child is set to new_child, while the prior definition is saved away
123        for later, when unset_all() is called.
124
125        This method supports the case where child_name is a staticmethod or a
126        classmethod of parent.
127        """
128        old_child = getattr(parent, child_name)
129
130        old_attribute = parent.__dict__.get(child_name)
131        if old_attribute is not None:
132            if isinstance(old_attribute, staticmethod):
133                old_child = staticmethod(old_child)
134            elif isinstance(old_attribute, classmethod):
135                old_child = classmethod(old_child.__func__)
136
137        self.cache.append((parent, old_child, child_name))
138        parent.__dict__[child_name] = new_child
139
140    def unset_all(self):
141        """Reverses all the Set() calls.
142
143        Restores things to their original definition. Its okay to call
144        unset_all() repeatedly, as later calls have no effect if no Set()
145        calls have been made.
146        """
147        # Undo calls to set() in reverse order, in case set() was called on the
148        # same arguments repeatedly (want the original call to be last one
149        # undone)
150        self.cache.reverse()
151
152        for parent, old_child, child_name in self.cache:
153            parent.__dict__[child_name] = old_child
154        self.cache = []
155