# HG changeset patch
# User Pierre-Yves David <pierre-yves.david@fb.com>
# Date 2014-07-31 02:04:50
# Node ID 7986e99bb69aa4562e5d42a63ce85117c431de4d
# Parent  7d976d71684c6d03bd348ca72607caca40ad3ca7

push: rework the bundle2partsgenerators logic

Instead of a single list of functions, we now have a list of names and
a mapping of names to functions. This simplifies wrapping of steps
from extensions. In the same move, declaration becomes decorator-based
(syntax sugar, nom nom nom!).

diff --git a/mercurial/exchange.py b/mercurial/exchange.py
--- a/mercurial/exchange.py
+++ b/mercurial/exchange.py
@@ -243,6 +243,31 @@ def _pushcheckoutgoing(pushop):
                              newbm)
     return True
 
+# List of names of steps to perform for an outgoing bundle2, order matters.
+b2partsgenorder = []
+
+# Mapping between step name and function
+#
+# This exists to help extensions wrap steps if necessary
+b2partsgenmapping = {}
+
+def b2partsgenerator(stepname):
+    """decorator for function generating bundle2 part
+
+    The function is added to the step -> function mapping and appended to the
+    list of steps.  Beware that decorated functions will be added in order
+    (this may matter).
+
+    You can only use this decorator for new steps, if you want to wrap a step
+    from an extension, attack the b2partsgenmapping dictionary directly."""
+    def dec(func):
+        assert stepname not in b2partsgenmapping
+        b2partsgenmapping[stepname] = func
+        b2partsgenorder.append(stepname)
+        return func
+    return dec
+
+@b2partsgenerator('changeset')
 def _pushb2ctx(pushop, bundler):
     """handle changegroup push through bundle2
 
@@ -269,8 +294,6 @@ def _pushb2ctx(pushop, bundler):
         pushop.ret = cgreplies['changegroup'][0]['return']
     return handlereply
 
-# list of function that may decide to add parts to an outgoing bundle2
-bundle2partsgenerators = [_pushb2ctx]
 
 def _pushbundle2(pushop):
     """push data to the remote using bundle2
@@ -282,7 +305,8 @@ def _pushbundle2(pushop):
     capsblob = bundle2.encodecaps(pushop.repo.bundle2caps)
     bundler.newpart('b2x:replycaps', data=capsblob)
     replyhandlers = []
-    for partgen in bundle2partsgenerators:
+    for partgenname in b2partsgenorder:
+        partgen = b2partsgenmapping[partgenname]
         ret = partgen(pushop, bundler)
         if callable(ret):
             replyhandlers.append(ret)
diff --git a/tests/test-bundle2.t b/tests/test-bundle2.t
--- a/tests/test-bundle2.t
+++ b/tests/test-bundle2.t
@@ -964,7 +964,8 @@ Setting up
   >     raise util.Abort('Abandon ship!', hint="don't panic")
   > 
   > def uisetup(ui):
-  >     exchange.bundle2partsgenerators.insert(0, _pushbundle2failpart)
+  >     exchange.b2partsgenmapping['failpart'] = _pushbundle2failpart
+  >     exchange.b2partsgenorder.insert(0, 'failpart')
   > 
   > EOF