# HG changeset patch
# User Martin von Zweigbergk <martinvonz@google.com>
# Date 1503359165 25200
# Node ID 47e52f079a579b652deeae2bef6b71d7f17232a6
# Parent  da07367d683b32ac08c61f704a42f9df82bcda89
extensions: add wrappedfunction() context manager

Several extensions exist that temporarily want to wrap a function (at
least narrowhg, any many of the extensions in hg-experimental). That's
why we have the unwrapfunction() that was introduced in 19578bb84731
(extensions: add unwrapfunction to undo wrapfunction, 2016-08-10).

This patch adds a simple wrappedfunction() that returns a context
manager.

Differential Revision: https://phab.mercurial-scm.org/D472

diff -r da07367d683b -r 47e52f079a57 mercurial/extensions.py
--- a/mercurial/extensions.py	Tue Aug 29 18:20:50 2017 -0700
+++ b/mercurial/extensions.py	Mon Aug 21 16:46:05 2017 -0700
@@ -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, self._wrapper)
+
 def wrapfunction(container, funcname, wrapper):
     '''Wrap the function named funcname in container
 
diff -r da07367d683b -r 47e52f079a57 tests/test-extensions-wrapfunction.py
--- a/tests/test-extensions-wrapfunction.py	Tue Aug 29 18:20:50 2017 -0700
+++ b/tests/test-extensions-wrapfunction.py	Mon Aug 21 16:46:05 2017 -0700
@@ -37,3 +37,20 @@
 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())
+        # Bad programmer forgets to unwrap the function, but the context
+        # managers still unwrap their wrappings.
+        extensions.wrapfunction(dummy, 'getstack', wrappers[2])
+        print('context manager', dummy.getstack())
+    print('context manager', dummy.getstack())
+print('context manager', dummy.getstack())
diff -r da07367d683b -r 47e52f079a57 tests/test-extensions-wrapfunction.py.out
--- a/tests/test-extensions-wrapfunction.py.out	Tue Aug 29 18:20:50 2017 -0700
+++ b/tests/test-extensions-wrapfunction.py.out	Mon Aug 21 16:46:05 2017 -0700
@@ -12,3 +12,9 @@
 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 [2, 0, 1, 'orig']
+context manager [2, 1, 'orig']
+context manager [2, 'orig']