diff --git a/mercurial/extensions.py b/mercurial/extensions.py --- a/mercurial/extensions.py +++ b/mercurial/extensions.py @@ -399,6 +399,21 @@ raise AttributeError(r"type '%s' has no property '%s'" % ( cls, propname)) +class wrappedfunction(object): + '''context manager for temporarily wrapping a function''' + + def __init__(self, container, funcname, wrapper): + assert callable(wrapper) + self._container = container + self._funcname = funcname + self._wrapper = wrapper + + def __enter__(self): + wrapfunction(self._container, self._funcname, self._wrapper) + + def __exit__(self, exctype, excvalue, traceback): + unwrapfunction(self._container, self._funcname) + def wrapfunction(container, funcname, wrapper): '''Wrap the function named funcname in container diff --git a/tests/test-extensions-wrapfunction.py b/tests/test-extensions-wrapfunction.py --- a/tests/test-extensions-wrapfunction.py +++ b/tests/test-extensions-wrapfunction.py @@ -37,3 +37,16 @@ batchwrap(wrappers + [wrappers[0]]) batchunwrap([(wrappers[i] if i >= 0 else None) for i in [3, None, 0, 4, 0, 2, 1, None]]) + +wrap0 = extensions.wrappedfunction(dummy, 'getstack', wrappers[0]) +wrap1 = extensions.wrappedfunction(dummy, 'getstack', wrappers[1]) + +# Use them in a different order from how they were created to check that +# the wrapping happens in __enter__, not in __init__ +print('context manager', dummy.getstack()) +with wrap1: + print('context manager', dummy.getstack()) + with wrap0: + print('context manager', dummy.getstack()) + print('context manager', dummy.getstack()) +print('context manager', dummy.getstack()) diff --git a/tests/test-extensions-wrapfunction.py.out b/tests/test-extensions-wrapfunction.py.out --- a/tests/test-extensions-wrapfunction.py.out +++ b/tests/test-extensions-wrapfunction.py.out @@ -12,3 +12,8 @@ unwrap 2: 2: [1, 'orig'] unwrap 1: 1: ['orig'] unwrap -: -: IndexError +context manager ['orig'] +context manager [1, 'orig'] +context manager [0, 1, 'orig'] +context manager [1, 'orig'] +context manager ['orig']