[PATCH 2 of 4] bundle2: separate bundle10 and bundle2 cases in getbundle()

[Mike Hommey]

# HG changeset patch
# User Mike Hommey <mh at glandium.org>
# Date 1411613277 -32400
#      Thu Sep 25 11:47:57 2014 +0900
# Node ID 7abbac9233058376fdcbfef0549e6714e3bf9711
# Parent  5c959b16e3f54b2cdc212264a5354574486e49b3
bundle2: separate bundle10 and bundle2 cases in getbundle()

The primary goal is to make it easier for extensions to alter how bundle2
parts are laid out. They now can use the getbundle2partsgenerator decorator
to add new parts, or directly act on getbundle2partsmapping to wrap existing
part functions.

Note the 'request for bundle10 must include changegroup' error was kept
under the same conditions as before, although the logic changes don't make
it obvious.

diff --git a/mercurial/exchange.py b/mercurial/exchange.py
--- a/mercurial/exchange.py
+++ b/mercurial/exchange.py
@@ -956,75 +956,117 @@ def _pullobsolete(pullop):
 
 def caps20to10(repo):
     """return a set with appropriate options to use bundle20 during getbundle"""
     caps = set(['HG2X'])
     capsblob = bundle2.encodecaps(bundle2.getrepocaps(repo))
     caps.add('bundle2=' + urllib.quote(capsblob))
     return caps
 
+# List of names of steps to perform for a bundle2 for getbundle, order matters.
+getbundle2partsorder = []
+
+# Mapping between step name and function
+#
+# This exists to help extensions wrap steps if necessary
+getbundle2partsmapping = {}
+
+def getbundle2partsgenerator(stepname):
+    """decorator for function generating bundle2 part for getbundle
+
+    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 getbundle2partsmapping dictionary directly."""
+    def dec(func):
+        assert stepname not in getbundle2partsmapping
+        getbundle2partsmapping[stepname] = func
+        getbundle2partsorder.append(stepname)
+        return func
+    return dec
+
 def getbundle(repo, source, heads=None, common=None, bundlecaps=None,
               **kwargs):
     """return a full bundle (with potentially multiple kind of parts)
 
     Could be a bundle HG10 or a bundle HG2X depending on bundlecaps
     passed. For now, the bundle can contain only changegroup, but this will
     changes when more part type will be available for bundle2.
 
     This is different from changegroup.getchangegroup that only returns an HG10
     changegroup bundle. They may eventually get reunited in the future when we
     have a clearer idea of the API we what to query different data.
 
     The implementation is at a very early stage and will get massive rework
     when the API of bundle is refined.
     """
-    cg = None
-    if kwargs.get('cg', True):
-        # build changegroup bundle here.
-        cg = changegroup.getchangegroup(repo, source, heads=heads,
-                                         common=common, bundlecaps=bundlecaps)
-    elif 'HG2X' not in bundlecaps:
-        raise ValueError(_('request for bundle10 must include changegroup'))
+    # bundle10 case
     if bundlecaps is None or 'HG2X' not in bundlecaps:
+        if bundlecaps and not kwargs.get('cg', True):
+            raise ValueError(_('request for bundle10 must include changegroup'))
+
         if kwargs:
             raise ValueError(_('unsupported getbundle arguments: %s')
                              % ', '.join(sorted(kwargs.keys())))
-        return cg
-    # very crude first implementation,
-    # the bundle API will change and the generation will be done lazily.
+        return changegroup.getchangegroup(repo, source, heads=heads,
+                                          common=common, bundlecaps=bundlecaps)
+
+    # bundle20 case
     b2caps = {}
     for bcaps in bundlecaps:
         if bcaps.startswith('bundle2='):
             blob = urllib.unquote(bcaps[len('bundle2='):])
             b2caps.update(bundle2.decodecaps(blob))
     bundler = bundle2.bundle20(repo.ui, b2caps)
+
+    for name in getbundle2partsorder:
+        func = getbundle2partsmapping[name]
+        func(bundler, repo, source, heads=heads, common=common,
+             bundlecaps=bundlecaps, b2caps=b2caps, **kwargs)
+
+    return util.chunkbuffer(bundler.getchunks())
+
+ at getbundle2partsgenerator('changegroup')
+def _getbundlechangegrouppart(bundler, repo, source, heads=None, common=None,
+                              bundlecaps=None, b2caps=None, **kwargs):
+    """add a changegroup part to the requested bundle"""
+    cg = None
+    if kwargs.get('cg', True):
+        # build changegroup bundle here.
+        cg = changegroup.getchangegroup(repo, source, heads=heads,
+                                        common=common, bundlecaps=bundlecaps)
+
     if cg:
         bundler.newpart('b2x:changegroup', data=cg.getchunks())
+
+ at getbundle2partsgenerator('listkeys')
+def _getbundlelistkeysparts(bundler, repo, source, heads=None, common=None,
+                           bundlecaps=None, b2caps=None, **kwargs):
+    """add parts containing listkeys namespaces to the requested bundle"""
     listkeys = kwargs.get('listkeys', ())
     for namespace in listkeys:
         part = bundler.newpart('b2x:listkeys')
         part.addparam('namespace', namespace)
         keys = repo.listkeys(namespace).items()
         part.data = pushkey.encodekeys(keys)
-    _getbundleobsmarkerpart(bundler, repo, source, heads=heads, common=common,
-                            bundlecaps=bundlecaps, b2caps=b2caps, **kwargs)
-    _getbundleextrapart(bundler, repo, source, heads=heads, common=common,
-                        bundlecaps=bundlecaps, b2caps=b2caps, **kwargs)
-    return util.chunkbuffer(bundler.getchunks())
 
+ at getbundle2partsgenerator('obsmarkers')
 def _getbundleobsmarkerpart(bundler, repo, source, heads=None, common=None,
                             bundlecaps=None, b2caps=None, **kwargs):
     """add an obsolescence markers part to the requested bundle"""
     if kwargs.get('obsmarkers', False):
         if heads is None:
             heads = repo.heads()
         subset = [c.node() for c in repo.set('::%ln', heads)]
         markers = repo.obsstore.relevantmarkers(subset)
         buildobsmarkerspart(bundler, markers)
 
+ at getbundle2partsgenerator('extra')
 def _getbundleextrapart(bundler, repo, source, heads=None, common=None,
                         bundlecaps=None, b2caps=None, **kwargs):
     """hook function to let extensions add parts to the requested bundle"""
     pass
 
 def check_heads(repo, their_heads, context):
     """check if the heads of a repo have been modified