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

[Mike Hommey]

# HG changeset patch
# User Mike Hommey <mh at glandium.org>
# Date 1411539963 -32400
#      Wed Sep 24 15:26:03 2014 +0900
# Node ID 07e31801a4b3fb762bf3cecf9875186cc182dcd0
# Parent  bc01442c6e030eee52c5f5524a28a50b1c25a4a6
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. While I don't expect this to be the final hookable form,
this hookability is necessary for upcoming tests for new bundle2 features.

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
@@ -971,61 +971,85 @@ def getbundle(repo, source, heads=None, 
 
     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,
+        return changegroup.getchangegroup(repo, source, heads=heads,
+                                          common=common, bundlecaps=bundlecaps)
+
+    # bundle20 case ; very crude first implementation,
     # the bundle API will change and the generation will be done lazily.
     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)
+
+    # Use a separate function to ease hooking the whole bundle content.
+    _getbundle2parts(bundler, repo, source, heads=heads, common=common,
+                     bundlecaps=bundlecaps, b2caps=b2caps, **kwargs)
+
+    return util.chunkbuffer(bundler.getchunks())
+
+def _getbundle2parts(bundler, repo, source, heads=None, common=None,
+                     bundlecaps=None, b2caps=None, **kwargs):
+    partfuncs = (
+        _getbundlechangegrouppart,
+        _getbundlelistkeyspart,
+        _getbundleobsmarkerpart,
+        _getbundleextrapart,
+    )
+    for func in partfuncs:
+        func(bundler, repo, source, heads=heads, common=common,
+             bundlecaps=bundlecaps, b2caps=b2caps, **kwargs)
+
+def _getbundlechangegrouppart(bundler, repo, source, heads=None, common=None,
+                              bundlecaps=None, b2caps=None, **kwargs):
+    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())
+
+def _getbundlelistkeyspart(bundler, repo, source, heads=None, common=None,
+                           bundlecaps=None, b2caps=None, **kwargs):
     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, **kwargs)
-    _getbundleextrapart(bundler, repo, source, heads=heads, common=common,
-                        bundlecaps=bundlecaps, **kwargs)
-    return util.chunkbuffer(bundler.getchunks())
 
 def _getbundleobsmarkerpart(bundler, repo, source, heads=None, common=None,
-                            bundlecaps=None, **kwargs):
+                            bundlecaps=None, b2caps=None, **kwargs):
     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)
 
 def _getbundleextrapart(bundler, repo, source, heads=None, common=None,
-                        bundlecaps=None, **kwargs):
+                        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
 
     Used by peer for unbundling.
     """