[PATCH evolve-ext] exthelper: wrap bundle2 processing correctly

[Eric Sumner]

# HG changeset patch
# User Eric Sumner <ericsumner at fb.com>
# Date 1429121236 14400
#      Wed Apr 15 14:07:16 2015 -0400
# Branch stable
# Node ID 488998e6b79734bffa39e1db80086900ae4de2ae
# Parent  1e7c8046a9f461bda3c8f6003dae65e235af8bb6
exthelper: wrap bundle2 processing correctly

Wrapping the bundle2 functions doesn't do anything; we need to update the
dictionaries maintained by their decorators instead.  This patch uses a new
decorator to store the wrapped functions in the correct place.

diff --git a/hgext/evolve.py b/hgext/evolve.py
--- a/hgext/evolve.py
+++ b/hgext/evolve.py
@@ -118,6 +118,7 @@
         self._commandwrappers = []
         self._extcommandwrappers = []
         self._functionwrappers = []
+        self._dictwrappers = []
         self._duckpunchers = []
 
     def final_uisetup(self, ui):
@@ -145,6 +146,10 @@
             extensions.wrapcommand(commands.table, command, wrapper)
         for cont, funcname, wrapper in self._functionwrappers:
             extensions.wrapfunction(cont, funcname, wrapper)
+
+        for cont, key, wrapper in self._dictwrappers:
+            orig = cont[key]
+            cont[key] = extensions.bind(wrapper, orig)
         for c in self._uicallables:
             c(ui)
 
@@ -304,6 +309,16 @@
             return wrapper
         return dec
 
+    def wrapdict(self, container, key):
+        """Decorated function is a function wrapper
+
+        This performs the same action as wrapfunction, but the container is
+        a dictionary instead of an object with attributes"""
+        def dec(wrapper):
+            self._dictwrappers.append((container, key, wrapper))
+            return wrapper
+        return dec
+
     def addattr(self, container, funcname):
         """Decorated function is to be added to the container
 
@@ -2285,38 +2300,19 @@
         topic = 'OBSEXC'
     ui.progress(topic, *args, **kwargs)
 
-if getattr(exchange, '_pushdiscoveryobsmarkers', None) is not None:
-    @eh.wrapfunction(exchange, '_pushdiscoveryobsmarkers')
-    def _pushdiscoveryobsmarkers(orig, pushop):
-        if (obsolete._enabled
-            and pushop.repo.obsstore
-            and 'obsolete' in pushop.remote.listkeys('namespaces')):
-            repo = pushop.repo
-            obsexcmsg(repo.ui, "computing relevant nodes\n")
-            revs = list(repo.revs('::%ln', pushop.futureheads))
-            unfi = repo.unfiltered()
-            cl = unfi.changelog
-            if not pushop.remote.capable('_evoext_obshash_0'):
-                # do not trust core yet
-                # return orig(pushop)
-                nodes = [cl.node(r) for r in revs]
-                if nodes:
-                    obsexcmsg(repo.ui, "computing markers relevant to %i nodes\n"
-                                       % len(nodes))
-                    pushop.outobsmarkers = repo.obsstore.relevantmarkers(nodes)
-                else:
-                    obsexcmsg(repo.ui, "markers already in sync\n")
-                    pushop.outobsmarkers = []
-                    pushop.outobsmarkers = repo.obsstore.relevantmarkers(nodes)
-                return
-
-            common = []
-            obsexcmsg(repo.ui, "looking for common markers in %i nodes\n"
-                               % len(revs))
-            commonrevs = list(unfi.revs('::%ln', pushop.outgoing.commonheads))
-            common = findcommonobsmarkers(pushop.ui, unfi, pushop.remote, commonrevs)
-
-            revs = list(unfi.revs('%ld - (::%ln)', revs, common))
+ at eh.wrapdict(exchange.pushdiscoverymapping, 'obsmarker')
+def _pushdiscoveryobsmarkers(orig, pushop):
+    if (obsolete._enabled
+        and pushop.repo.obsstore
+        and 'obsolete' in pushop.remote.listkeys('namespaces')):
+        repo = pushop.repo
+        obsexcmsg(repo.ui, "computing relevant nodes\n")
+        revs = list(repo.revs('::%ln', pushop.futureheads))
+        unfi = repo.unfiltered()
+        cl = unfi.changelog
+        if not pushop.remote.capable('_evoext_obshash_0'):
+            # do not trust core yet
+            # return orig(pushop)
             nodes = [cl.node(r) for r in revs]
             if nodes:
                 obsexcmsg(repo.ui, "computing markers relevant to %i nodes\n"
@@ -2325,6 +2321,24 @@
             else:
                 obsexcmsg(repo.ui, "markers already in sync\n")
                 pushop.outobsmarkers = []
+                pushop.outobsmarkers = repo.obsstore.relevantmarkers(nodes)
+            return
+
+        common = []
+        obsexcmsg(repo.ui, "looking for common markers in %i nodes\n"
+                           % len(revs))
+        commonrevs = list(unfi.revs('::%ln', pushop.outgoing.commonheads))
+        common = findcommonobsmarkers(pushop.ui, unfi, pushop.remote, commonrevs)
+
+        revs = list(unfi.revs('%ld - (::%ln)', revs, common))
+        nodes = [cl.node(r) for r in revs]
+        if nodes:
+            obsexcmsg(repo.ui, "computing markers relevant to %i nodes\n"
+                               % len(nodes))
+            pushop.outobsmarkers = repo.obsstore.relevantmarkers(nodes)
+        else:
+            obsexcmsg(repo.ui, "markers already in sync\n")
+            pushop.outobsmarkers = []
 
 @eh.wrapfunction(wireproto, 'capabilities')
 def discocapabilities(orig, repo, proto):
@@ -2601,22 +2615,20 @@
             kwargs['evo_obscommon'] = common
     return ret
 
-if getattr(exchange, '_getbundleobsmarkerpart', None) is not None:
-    @eh.wrapfunction(exchange, '_getbundleobsmarkerpart')
-    def _getbundleobsmarkerpart(orig, bundler, repo, source, heads=None, common=None,
-                                bundlecaps=None, **kwargs):
-        if 'evo_obscommon' not in kwargs:
-            return orig(bundler, repo, source, heads, common, bundlecaps, **kwargs)
-
-        if kwargs.get('obsmarkers', False):
-            if heads is None:
-                heads = repo.heads()
-            obscommon = kwargs.get('evo_obscommon', ())
-            obsset = repo.set('::%ln - ::%ln', heads, obscommon)
-            subset = [c.node() for c in obsset]
-            markers = repo.obsstore.relevantmarkers(subset)
-            exchange.buildobsmarkerspart(bundler, markers)
-
+ at eh.wrapdict(exchange.getbundle2partsmapping, 'obsmarkers')
+def _getbundleobsmarkerpart(orig, bundler, repo, source, heads=None, common=None,
+                        bundlecaps=None, **kwargs):
+    if 'evo_obscommon' not in kwargs:
+        return orig(bundler, repo, source, heads, common, bundlecaps, **kwargs)
+
+    if kwargs.get('obsmarkers', False):
+        if heads is None:
+            heads = repo.heads()
+        obscommon = kwargs.get('evo_obscommon', ())
+        obsset = repo.set('::%ln - ::%ln', heads, obscommon)
+        subset = [c.node() for c in obsset]
+        markers = repo.obsstore.relevantmarkers(subset)
+        exchange.buildobsmarkerspart(bundler, markers)
 
 @eh.wrapfunction(exchange, '_pullobsolete')
 def _pullobsolete(orig, pullop):