diff options
-rw-r--r-- | alembic/revision.py | 70 | ||||
-rw-r--r-- | tests/test_revision.py | 20 |
2 files changed, 60 insertions, 30 deletions
diff --git a/alembic/revision.py b/alembic/revision.py index c7dabb8..4f2297b 100644 --- a/alembic/revision.py +++ b/alembic/revision.py @@ -505,6 +505,7 @@ class RevisionMap(object): else: lowers = requested_lowers + # represents all nodes we will produce total_space = set( rev.revision for rev in upper_ancestors).intersection( rev.revision for rev @@ -514,43 +515,52 @@ class RevisionMap(object): if not total_space: raise RangeNotAncestorError(lower, upper) - branch_endpoints = set( - rev.revision for rev in + # organize branch points to be consumed separately from + # member nodes + branch_todo = set( + rev for rev in (self._revision_map[rev] for rev in total_space) if rev.is_branch_point and len(total_space.intersection(rev.nextrev)) > 1 ) + # it's not possible for any "uppers" to be in branch_todo, + # because the .nextrev of those nodes is not in total_space + #assert not branch_todo.intersection(uppers) + todo = collections.deque( r for r in uppers if r.revision in total_space) - stop = set(lowers) - if implicit_base: - stop.update(self._get_ancestor_nodes(requested_lowers)) - while todo: - stop.update( - rev.revision for rev in todo - if rev.revision in branch_endpoints) - rev = todo.popleft() - - # do depth first for elements within the branches - todo.extendleft([ - self._revision_map[downrev] - for downrev in reversed(rev._down_revision_tuple) - if downrev not in branch_endpoints and downrev not in stop - and downrev in total_space]) - - # then put the actual branch points at the end of the - # list for subsequent traversal - todo.extend([ - self._revision_map[downrev] - for downrev in rev._down_revision_tuple - if downrev in branch_endpoints and downrev not in stop - and downrev in total_space - ]) - - if not inclusive and rev in requested_lowers: - continue - yield rev + + # iterate for total_space being emptied out + while total_space: + # when everything non-branch pending is consumed, + # add to the todo any branch nodes that have no + # descendants left in the queue + if not todo: + todo.extendleft( + rev for rev in branch_todo + if not rev.nextrev.intersection(total_space) + ) + branch_todo.difference_update(todo) + + # iterate nodes that are in the immediate todo + while todo: + rev = todo.popleft() + total_space.remove(rev.revision) + + # do depth first for elements within branches, + # don't consume any actual branch nodes + todo.extendleft([ + self._revision_map[downrev] + for downrev in reversed(rev._down_revision_tuple) + if self._revision_map[downrev] not in branch_todo + and downrev in total_space]) + + if not inclusive and rev in requested_lowers: + continue + yield rev + + assert not branch_todo class Revision(object): diff --git a/tests/test_revision.py b/tests/test_revision.py index f5a9226..10fae91 100644 --- a/tests/test_revision.py +++ b/tests/test_revision.py @@ -319,6 +319,26 @@ class LabeledBranchTest(DownIterateTest): ) +class LongShortBranchTest(DownIterateTest): + def setUp(self): + self.map = RevisionMap( + lambda: [ + Revision('a', ()), + Revision('b1', ('a',)), + Revision('b2', ('a',)), + Revision('c1', ('b1',)), + Revision('d11', ('c1',)), + Revision('d12', ('c1',)), + ] + ) + + def test_iterate_full(self): + self._assert_iteration( + "heads", "base", + ['b2', 'd11', 'd12', 'c1', 'b1', 'a'] + ) + + class MultipleBranchTest(DownIterateTest): def setUp(self): self.map = RevisionMap( |