diff --git a/mercurial/extensions.py b/mercurial/extensions.py --- a/mercurial/extensions.py +++ b/mercurial/extensions.py @@ -399,6 +399,34 @@ 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) + + origfn = getattr(container, funcname) + assert callable(origfn) + wrap = bind(wrapper, origfn) + _updatewrapper(wrap, origfn, wrapper) + setattr(container, funcname, wrap) + + self._container = container + self._funcname = funcname + self._origfn = origfn + + def __enter__(self): + return self + + def __exit__(self, exctype, excvalue, traceback): + unwrapfunction(self._container, self._funcname) + + def __get__(self, instance, owner): + util.nouideprecwarn("Using the return value of wrapfunction() as a " + "function is deprecated. Use it as context manager " + "instead","4.4", stacklevel=2) + return bind(self._origfn, instance) + def wrapfunction(container, funcname, wrapper): '''Wrap the function named funcname in container @@ -432,14 +460,7 @@ your end users, you should play nicely with others by using the subclass trick. ''' - assert callable(wrapper) - - origfn = getattr(container, funcname) - assert callable(origfn) - wrap = bind(wrapper, origfn) - _updatewrapper(wrap, origfn, wrapper) - setattr(container, funcname, wrap) - return origfn + return _wrappedfunction(container, funcname, wrapper) def unwrapfunction(container, funcname, wrapper=None): '''undo wrapfunction 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,11 @@ batchwrap(wrappers + [wrappers[0]]) batchunwrap([(wrappers[i] if i >= 0 else None) for i in [3, None, 0, 4, 0, 2, 1, None]]) + +print('context manager', dummy.getstack()) +with extensions.wrapfunction(dummy, 'getstack', wrappers[0]): + print('context manager', dummy.getstack()) + with extensions.wrapfunction(dummy, 'getstack', wrappers[1]): + 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 [0, 'orig'] +context manager [1, 0, 'orig'] +context manager [0, 'orig'] +context manager ['orig'] diff --git a/tests/test-filecache.py b/tests/test-filecache.py --- a/tests/test-filecache.py +++ b/tests/test-filecache.py @@ -129,20 +129,16 @@ def wrapinit(orig, *args, **kwargs): pass - originit = extensions.wrapfunction(util.cachestat, '__init__', wrapinit) - origcacheable = extensions.wrapfunction(util.cachestat, 'cacheable', - wrapcacheable) + with extensions.wrapfunction(util.cachestat, '__init__', wrapinit),\ + extensions.wrapfunction(util.cachestat, 'cacheable', wrapcacheable): - for fn in ['x', 'y']: - try: - os.remove(fn) - except OSError: - pass + for fn in ['x', 'y']: + try: + os.remove(fn) + except OSError: + pass - basic(fakerepo()) - - util.cachestat.cacheable = origcacheable - util.cachestat.__init__ = originit + basic(fakerepo()) def test_filecache_synced(): # test old behavior that caused filecached properties to go out of sync