diff --git a/mercurial/bundle2.py b/mercurial/bundle2.py --- a/mercurial/bundle2.py +++ b/mercurial/bundle2.py @@ -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