bundle2: separate bundle10 and bundle2 cases in getbundle()
authorMike Hommey <mh@glandium.org>
Thu, 25 Sep 2014 11:47:57 +0900
changeset 22542 6b180a0c703e
parent 22541 4e1a80c022a4
child 22543 6a669d5a62b8
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.
mercurial/exchange.py
--- a/mercurial/exchange.py	Thu Sep 25 11:11:37 2014 +0900
+++ b/mercurial/exchange.py	Thu Sep 25 11:47:57 2014 +0900
@@ -961,6 +961,30 @@
     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)
@@ -976,40 +1000,57 @@
     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())
+
+@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())
+
+@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())
 
+@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"""
@@ -1020,6 +1061,7 @@
         markers = repo.obsstore.relevantmarkers(subset)
         buildobsmarkerspart(bundler, markers)
 
+@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"""