From patchwork Fri Jul  1 12:09:05 2016
Content-Type: text/plain; charset="utf-8"
MIME-Version: 1.0
Content-Transfer-Encoding: 7bit
Subject: [3,of,3] extensions: add unwrapfunction to undo wrapfunction
From: Jun Wu <quark@fb.com>
X-Patchwork-Id: 15700
Message-Id: <7a68b01000d12e1795d1.1467374945@x1c>
To: <mercurial-devel@mercurial-scm.org>
Date: Fri, 1 Jul 2016 13:09:05 +0100

# HG changeset patch
# User Jun Wu <quark@fb.com>
# Date 1467374153 -3600
#      Fri Jul 01 12:55:53 2016 +0100
# Node ID 7a68b01000d12e1795d142cdb32fdd62a5ffc8cb
# Parent  de798903374922317eb2cbca9733ba0cea415780
# Available At https://bitbucket.org/quark-zju/hg-draft
#              hg pull https://bitbucket.org/quark-zju/hg-draft -r 7a68b01000d1
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. And it allows extensions running with chg to add and remove their
wrappers dynamically in reposetup without starting a new chgserver.

diff --git a/mercurial/extensions.py b/mercurial/extensions.py
--- a/mercurial/extensions.py
+++ b/mercurial/extensions.py
@@ -303,6 +303,19 @@ def wrapfunction(container, funcname, wr
     setattr(container, funcname, wrap)
     return origfn
 
+def unwrapfunction(container, funcname, wrapper):
+    '''Undo wrapfunction
+
+    Pass the same arguments with wrapfunction to undo a wrap.
+    Handles stacked wrappers correctly.
+    '''
+    chain = getwrapperchain(container, funcname)
+    origfn = chain.pop()
+    chain.remove(wrapper)
+    setattr(container, funcname, origfn)
+    for wrapper in reversed(chain):
+        wrapfunction(container, funcname, wrapper)
+
 def getwrapperchain(container, funcname):
     '''get a chain of wrappers of a function
 
diff --git a/tests/test-extensions-wrapfunction.py b/tests/test-extensions-wrapfunction.py
new file mode 100644
--- /dev/null
+++ b/tests/test-extensions-wrapfunction.py
@@ -0,0 +1,24 @@
+from __future__ import absolute_import, print_function
+
+from mercurial import extensions
+
+def genwrap(x):
+    def f(orig, *args, **kwds):
+        return [x] + orig(*args, **kwds)
+    return f
+
+wrappers = [genwrap(i) for i in range(5)]
+
+class dummyclass(object):
+    def foo(self):
+        return ['orig']
+
+dummy = dummyclass()
+
+for w in wrappers + [wrappers[0]]:
+    extensions.wrapfunction(dummy, 'foo', w)
+    print(dummy.foo())
+
+for i in [3, 0, 4, 0, 2, 1]:
+    extensions.unwrapfunction(dummy, 'foo', wrappers[i])
+    print(dummy.foo())
diff --git a/tests/test-extensions-wrapfunction.py.out b/tests/test-extensions-wrapfunction.py.out
new file mode 100644
--- /dev/null
+++ b/tests/test-extensions-wrapfunction.py.out
@@ -0,0 +1,12 @@
+[0, 'orig']
+[1, 0, 'orig']
+[2, 1, 0, 'orig']
+[3, 2, 1, 0, 'orig']
+[4, 3, 2, 1, 0, 'orig']
+[0, 4, 3, 2, 1, 0, 'orig']
+[0, 4, 2, 1, 0, 'orig']
+[4, 2, 1, 0, 'orig']
+[2, 1, 0, 'orig']
+[2, 1, 'orig']
+[1, 'orig']
+['orig']