changeset 29777:19578bb84731

extensions: add unwrapfunction to undo wrapfunction Before this patch, we don't have a safe way to undo a wrapfunction because other extensions may wrap the same function and calling setattr will undo them accidentally. This patch adds an "unwrapfunction" to address the issue. It removes the wrapper from the wrapper chain, and re-wraps everything, which is not the most efficient but short and easy to understand. We can revisit the code if we have perf issues with long chains. The "undo" feature is useful in cases like wrapping a function just in a scope. Like, having a "select" command to interactively (using arrow keys) select content from some output (ex. smartlog). It could wrap "ui.label" to extract interesting texts just in the "select" command.
author Jun Wu <quark@fb.com>
date Wed, 10 Aug 2016 16:27:33 +0100
parents 8bf97c4c6c2a
children e5b794063fd4
files mercurial/extensions.py tests/test-extensions-wrapfunction.py tests/test-extensions-wrapfunction.py.out
diffstat 3 files changed, 73 insertions(+), 0 deletions(-) [+]
line wrap: on
line diff
--- a/mercurial/extensions.py	Wed Aug 10 15:21:42 2016 +0100
+++ b/mercurial/extensions.py	Wed Aug 10 16:27:33 2016 +0100
@@ -309,6 +309,26 @@
     setattr(container, funcname, wrap)
     return origfn
 
+def unwrapfunction(container, funcname, wrapper=None):
+    '''undo wrapfunction
+
+    If wrappers is None, undo the last wrap. Otherwise removes the wrapper
+    from the chain of wrappers.
+
+    Return the removed wrapper.
+    Raise IndexError if wrapper is None and nothing to unwrap; ValueError if
+    wrapper is not None but is not found in the wrapper chain.
+    '''
+    chain = getwrapperchain(container, funcname)
+    origfn = chain.pop()
+    if wrapper is None:
+        wrapper = chain[0]
+    chain.remove(wrapper)
+    setattr(container, funcname, origfn)
+    for w in reversed(chain):
+        wrapfunction(container, funcname, w)
+    return wrapper
+
 def getwrapperchain(container, funcname):
     '''get a chain of wrappers of a function
 
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/tests/test-extensions-wrapfunction.py	Wed Aug 10 16:27:33 2016 +0100
@@ -0,0 +1,39 @@
+from __future__ import absolute_import, print_function
+
+from mercurial import extensions
+
+def genwrapper(x):
+    def f(orig, *args, **kwds):
+        return [x] + orig(*args, **kwds)
+    f.x = x
+    return f
+
+def getid(wrapper):
+    return getattr(wrapper, 'x', '-')
+
+wrappers = [genwrapper(i) for i in range(5)]
+
+class dummyclass(object):
+    def getstack(self):
+        return ['orig']
+
+dummy = dummyclass()
+
+def batchwrap(wrappers):
+    for w in wrappers:
+        extensions.wrapfunction(dummy, 'getstack', w)
+        print('wrap %d: %s' % (getid(w), dummy.getstack()))
+
+def batchunwrap(wrappers):
+    for w in wrappers:
+        result = None
+        try:
+            result = extensions.unwrapfunction(dummy, 'getstack', w)
+            msg = str(dummy.getstack())
+        except (ValueError, IndexError) as e:
+            msg = e.__class__.__name__
+        print('unwrap %s: %s: %s' % (getid(w), getid(result), msg))
+
+batchwrap(wrappers + [wrappers[0]])
+batchunwrap([(wrappers[i] if i >= 0 else None)
+             for i in [3, None, 0, 4, 0, 2, 1, None]])
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/tests/test-extensions-wrapfunction.py.out	Wed Aug 10 16:27:33 2016 +0100
@@ -0,0 +1,14 @@
+wrap 0: [0, 'orig']
+wrap 1: [1, 0, 'orig']
+wrap 2: [2, 1, 0, 'orig']
+wrap 3: [3, 2, 1, 0, 'orig']
+wrap 4: [4, 3, 2, 1, 0, 'orig']
+wrap 0: [0, 4, 3, 2, 1, 0, 'orig']
+unwrap 3: 3: [0, 4, 2, 1, 0, 'orig']
+unwrap -: 0: [4, 2, 1, 0, 'orig']
+unwrap 0: 0: [4, 2, 1, 'orig']
+unwrap 4: 4: [2, 1, 'orig']
+unwrap 0: -: ValueError
+unwrap 2: 2: [1, 'orig']
+unwrap 1: 1: ['orig']
+unwrap -: -: IndexError