[PATCH 4 of 5 V2] bundle2: records processing results in the bundleoperation object

[pierre-yves.david at ens-lyon.org]

# HG changeset patch
# User Pierre-Yves David <pierre-yves.david at fb.com>
# Date 1396503470 25200
#      Wed Apr 02 22:37:50 2014 -0700
# Node ID 3583b40c588e443cea2501f59e4fe30f62d6ba8e
# Parent  10cf49dd923a8522df571f3619a0e2dfdeeeeb28
bundle2: records processing results in the bundleoperation object

Part handlers can now add records to the `bundleoperation` object This can be
used to help other part or to let the caller of the unbundling process to react
to the results.

diff --git a/mercurial/bundle2.py b/mercurial/bundle2.py
--- a/mercurial/bundle2.py
+++ b/mercurial/bundle2.py
@@ -181,10 +181,49 @@ def parthandler(parttype):
         assert lparttype not in parthandlermapping
         parthandlermapping[lparttype] = func
         return func
     return _decorator
 
+class unbundlerecords(object):
+    """keep record of what happen during and unbundle
+
+    New record are added using `records.add('cat', obj)`. Where 'cat' is a
+    category of record and obj is an arbitraty object.
+
+    `records['cat']` will return all entries of this category 'cat'.
+
+    Iterating on the object itself will yield `('category', obj)` tuple for all
+    entries.
+
+    All iterations happens in chronological order.
+    """
+
+    def __init__(self):
+        self._categories = {}
+        self._sequences = []
+
+    def add(self, category, entry):
+        """add a new record of a given category.
+
+        The entry can then be retrieved in the list returned by
+        self['category']."""
+        self._categories.setdefault(category, []).append(entry)
+        self._sequences.append((category, entry))
+
+    def __getitem__(self, cat):
+        return tuple(self._categories.get(cat, ()))
+
+    def __iter__(self):
+        return iter(self._sequences)
+
+    def __len__(self):
+        return len(self._sequences)
+
+    def __nonzero__(self):
+        return bool(self._sequences)
+
+
 class bundleoperation(object):
     """an object that represent a single bundle processing
 
     It purpose is to carry unbundle related objects and states.
 
@@ -200,10 +239,11 @@ class bundleoperation(object):
     """
 
     def __init__(self, repo):
         self.repo = repo
         self.ui = repo.ui
+        self.records = unbundlerecords()
 
 
 def processbundle(repo, unbundler):
     """This function process a bundle, apply effect to/from a repo
 
@@ -241,10 +281,11 @@ def processbundle(repo, unbundler):
             handler(op, part)
     except Exception:
         for part in iterparts:
             pass # consume the bundle content
         raise
+    return op
 
 class bundle20(object):
     """represent an outgoing bundle2 container
 
     Use the `addparam` method to add stream level parameter. and `addpart` to
diff --git a/tests/test-bundle2.t b/tests/test-bundle2.t
--- a/tests/test-bundle2.t
+++ b/tests/test-bundle2.t
@@ -22,12 +22,15 @@ Create an extension to test bundle2 API
   > 
   > @bundle2.parthandler('test:song')
   > def songhandler(op, part):
   >     """handle a "test:song" bundle2 part, printing the lyrics on stdin"""
   >     op.ui.write('The choir start singing:\n')
+  >     verses = 0
   >     for line in part.data.split('\n'):
   >         op.ui.write('    %s\n' % line)
+  >         verses += 1
+  >     op.records.add('song', {'verses': verses})
   > 
   > @command('bundle2',
   >          [('', 'param', [], 'stream level parameter'),
   >           ('', 'unknown', False, 'include an unknown mandatory part in the bundle'),
   >           ('', 'parts', False, 'include some arbitrary parts to the bundle'),],
@@ -73,17 +76,20 @@ Create an extension to test bundle2 API
   >     """process a bundle2 stream from stdin on the current repo"""
   >     try:
   >         lock = repo.lock()
   >         try:
   >             unbundler = bundle2.unbundle20(ui, sys.stdin)
-  >             bundle2.processbundle(repo, unbundler)
+  >             op = bundle2.processbundle(repo, unbundler)
   >         except KeyError, exc:
   >             raise util.Abort('missing support for %s' % exc)
   >     finally:
   >         lock.release()
   >         remains = sys.stdin.read()
   >         ui.write('%i unread bytes\n' % len(remains))
+  >     if op.records['song']:
+  >         totalverses = sum(r['verses'] for r in op.records['song'])
+  >         ui.write('%i total verses sung\n' % totalverses)
   > 
   > @command('statbundle2', [], '')
   > def cmdstatbundle2(ui, repo):
   >     """print statistic on the bundle2 container read from stdin"""
   >     unbundler = bundle2.unbundle20(ui, sys.stdin)
@@ -391,10 +397,11 @@ Process the bundle
   payload chunk size: 0
   ignoring unknown advisory part 'test:math'
   part header size: 0
   end of bundle2 stream
   0 unread bytes
+  3 total verses sung
 
 
   $ hg bundle2 --parts --unknown ../unknown.hg2
 
   $ hg unbundle2 < ../unknown.hg2