changeset 34150:e9e0e1143fc5

bundle2: move part iterator a separate class

Currently, the part iterator logic is tightly coupled with the part handling
logic, which means it's hard to replace the part handling logic without
duplicating the part iterator bits.

In a future diff we'll want to be able to replace all part handling, so let's
begin refactoring the part iterator logic to it's own class.

Differential Revision: https://phab.mercurial-scm.org/D703

--- a/mercurial/bundle2.py	Mon Sep 11 13:39:22 2017 -0700
+++ b/mercurial/bundle2.py	Wed Sep 13 17:16:45 2017 -0700
@@ -347,6 +347,16 @@
         _processchangegroup(op, unbundler, tr, source, url, **kwargs)
         return op
 
+class partiterator(object):
+    def __init__(self, unbundler):
+        self.unbundler = unbundler
+
+    def __enter__(self):
+        return enumerate(self.unbundler.iterparts())
+
+    def __exit__(self, type, value, tb):
+        pass
+
 def processbundle(repo, unbundler, transactiongetter=None, op=None):
     """This function process a bundle, apply effect to/from a repo
 
@@ -378,48 +388,49 @@
             msg.append(' with-transaction')
         msg.append('\n')
         repo.ui.debug(''.join(msg))
-    iterparts = enumerate(unbundler.iterparts())
-    part = None
-    nbpart = 0
-    try:
-        for nbpart, part in iterparts:
-            _processpart(op, part)
-    except Exception as exc:
-        # Any exceptions seeking to the end of the bundle at this point are
-        # almost certainly related to the underlying stream being bad.
-        # And, chances are that the exception we're handling is related to
-        # getting in that bad state. So, we swallow the seeking error and
-        # re-raise the original error.
-        seekerror = False
+
+    with partiterator(unbundler) as parts:
+        part = None
+        nbpart = 0
         try:
-            for nbpart, part in iterparts:
-                # consume the bundle content
-                part.seek(0, 2)
-        except Exception:
-            seekerror = True
+            for nbpart, part in parts:
+                _processpart(op, part)
+        except Exception as exc:
+            # Any exceptions seeking to the end of the bundle at this point are
+            # almost certainly related to the underlying stream being bad.
+            # And, chances are that the exception we're handling is related to
+            # getting in that bad state. So, we swallow the seeking error and
+            # re-raise the original error.
+            seekerror = False
+            try:
+                for nbpart, part in parts:
+                    # consume the bundle content
+                    part.seek(0, 2)
+            except Exception:
+                seekerror = True
 
-        # Small hack to let caller code distinguish exceptions from bundle2
-        # processing from processing the old format. This is mostly
-        # needed to handle different return codes to unbundle according to the
-        # type of bundle. We should probably clean up or drop this return code
-        # craziness in a future version.
-        exc.duringunbundle2 = True
-        salvaged = []
-        replycaps = None
-        if op.reply is not None:
-            salvaged = op.reply.salvageoutput()
-            replycaps = op.reply.capabilities
-        exc._replycaps = replycaps
-        exc._bundle2salvagedoutput = salvaged
+            # Small hack to let caller code distinguish exceptions from bundle2
+            # processing from processing the old format. This is mostly needed
+            # to handle different return codes to unbundle according to the type
+            # of bundle. We should probably clean up or drop this return code
+            # craziness in a future version.
+            exc.duringunbundle2 = True
+            salvaged = []
+            replycaps = None
+            if op.reply is not None:
+                salvaged = op.reply.salvageoutput()
+                replycaps = op.reply.capabilities
+            exc._replycaps = replycaps
+            exc._bundle2salvagedoutput = salvaged
 
-        # Re-raising from a variable loses the original stack. So only use
-        # that form if we need to.
-        if seekerror:
-            raise exc
-        else:
-            raise
-    finally:
-        repo.ui.debug('bundle2-input-bundle: %i parts total\n' % nbpart)
+            # Re-raising from a variable loses the original stack. So only use
+            # that form if we need to.
+            if seekerror:
+                raise exc
+            else:
+                raise
+        finally:
+            repo.ui.debug('bundle2-input-bundle: %i parts total\n' % nbpart)
 
     return op