# HG changeset patch
# User Durham Goode <durham@fb.com>
# Date 1505348205 25200
# Node ID e9e0e1143fc5e6b990e7521026996b23fc37ab79
# Parent  75cc1f1e11f2cd41e18b55842f5a7b19c31fa993
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

diff -r 75cc1f1e11f2 -r e9e0e1143fc5 mercurial/bundle2.py
--- 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