mercurial/bundle2.py
changeset 34150 e9e0e1143fc5
parent 34101 5ede882c249c
child 34151 550343626bb2
--- 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