# HG changeset patch
# User Pierre-Yves David <pierre-yves.david@fb.com>
# Date 2014-04-03 05:37:50
# Node ID 571f2903ff1e681a02de9e68ceb5cfcf9a6226f4
# Parent  329cd74b52bd9a2e72e179b5e49276c0c62b949e

bundle2: record processing results in the bundleoperation object

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

diff --git a/mercurial/bundle2.py b/mercurial/bundle2.py
--- a/mercurial/bundle2.py
+++ b/mercurial/bundle2.py
@@ -183,6 +183,44 @@ def parthandler(parttype):
         return func
     return _decorator
 
+class unbundlerecords(object):
+    """keep record of what happens during and unbundle
+
+    New records 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)` tuples
+    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 represents a single bundling process
 
@@ -202,6 +240,7 @@ 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
@@ -242,6 +281,7 @@ def processbundle(repo, unbundler):
         for part in iterparts:
             pass # consume the bundle content
         raise
+    return op
 
 class bundle20(object):
     """represent an outgoing bundle2 container
diff --git a/tests/test-bundle2.t b/tests/test-bundle2.t
--- a/tests/test-bundle2.t
+++ b/tests/test-bundle2.t
@@ -24,8 +24,11 @@ Create an extension to test bundle2 API
   > def songhandler(op, part):
   >     """handle a "test:song" bundle2 part, printing the lyrics on stdin"""
   >     op.ui.write('The choir starts 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'),
@@ -75,13 +78,16 @@ Create an extension to test bundle2 API
   >         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):
@@ -393,6 +399,7 @@ Process the bundle
   part header size: 0
   end of bundle2 stream
   0 unread bytes
+  3 total verses sung
 
 
   $ hg bundle2 --parts --unknown ../unknown.hg2