intercept.py 文件源码

def replaced(method, replacement, on=None):
    """A context manager which replaces the method passed in as `method` with the callable
       in `replacement` for the duration of the managed context.

       .. code-block:: python

          class SomethingElse(object):
              def foo(self, n, y='yyy'):
                  assert None, 'This should never be reached in this test'

          s = SomethingElse()

          def replacement(n, y=None):
              return y

          with replaced(s.foo, replacement):
              assert s.foo(1, y='a') == 'a'
              assert s.foo(2) == None

          assert s.foo(2) == 'yyy'
    """
    StubClass.signatures_match(method, replacement, ignore_self=True)

    is_unbound_method = six.PY2 and inspect.ismethod(method) and not method.im_self
    if (not inspect.ismethod(method) or is_unbound_method) and not on:
        raise ValueError('You have to supply a on= when stubbing an unbound method')
    target = on or six.get_method_self(method)

    if inspect.isfunction(method) or inspect.isbuiltin(method):
        method_name = method.__name__
    else:
        method_name = six.get_method_function(method).__name__
    saved_method = getattr(target, method_name)
    try:
        setattr(target, method_name, replacement)
        yield
    finally:
        setattr(target, method_name, saved_method)