From c6b56030ec3f7503046dcb5f87b3643f0858d08b Mon Sep 17 00:00:00 2001
From: Andrew Brown <brownan@gmail.com>
Date: Sun, 20 Nov 2011 10:50:37 -0500
Subject: [PATCH] dirtytile trees can now iterate at any level

Fixed some broken tests too.
---
 overviewer_core/quadtree.py | 33 ++++++++++++++++++++++-----------
 test/test_dirtytiles.py     | 17 +++++++++++++++--
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/overviewer_core/quadtree.py b/overviewer_core/quadtree.py
index 4664470..c7bd561 100644
--- a/overviewer_core/quadtree.py
+++ b/overviewer_core/quadtree.py
@@ -621,7 +621,8 @@ class DirtyTiles(object):
 
         """
         # Stores the depth of the tree according to this node. This is not the
-        # depth of this node, but rather the number of levels below this node.
+        # depth of this node, but rather the number of levels below this node
+        # (including this node).
         self.depth = depth
 
         # the self.children array holds the 4 children of this node. This
@@ -697,21 +698,31 @@ class DirtyTiles(object):
                     if all(x is True for x in self.children):
                         return True
 
-    def iterate_dirty(self, depth=None):
+    def iterate_dirty(self, level=None):
         """Returns an iterator over every dirty tile in this subtree. Each item
         yielded is a sequence of integers representing the quadtree path to the
         dirty tile. Yielded sequences are of length self.depth.
 
-        If zoom is None, iterates over tiles of the highest level, i.e.
-        worldtiles. If zoom is a value between 0 and the depth, iterates over
-        tiles at that zoom level. Zoom level 0 is zoomed all the way out, zoom
-        level `depth` is all the way in.
+        If level is None, iterates over tiles of the highest level, i.e.
+        worldtiles. If level is a value between 0 and the depth of this tree,
+        this method iterates over tiles at that level. Zoom level 0 is zoomed
+        all the way out, zoom level `depth` is all the way in.
+
+        In other words, specifying level causes the tree to be iterated as if
+        it was only that depth.
 
         """
-        return (tuple(reversed(rpath)) for rpath in self._iterate_dirty_helper())
+        if level is None:
+            todepth = 1
+        else:
+            if not (level > 0 and level <= self.depth):
+                raise ValueError("Level parameter must be between 1 and %s" % self.depth)
+            todepth = self.depth - level + 1
 
-    def _iterate_dirty_helper(self):
-        if self.depth == 1:
+        return (tuple(reversed(rpath)) for rpath in self._iterate_dirty_helper(todepth))
+
+    def _iterate_dirty_helper(self, todepth):
+        if self.depth == todepth:
             # Base case
             if self.children[0]: yield [0]
             if self.children[1]: yield [1]
@@ -723,13 +734,13 @@ class DirtyTiles(object):
             for c, child in enumerate(self.children):
                 if child == True:
                     # All dirty down this subtree, iterate over every leaf
-                    for x in iterate_base4(self.depth-1):
+                    for x in iterate_base4(self.depth-todepth):
                         x = list(x)
                         x.append(c)
                         yield x
                 elif child != False:
                     # Mixed dirty/clean down this subtree, recurse
-                    for path in child._iterate_dirty_helper():
+                    for path in child._iterate_dirty_helper(todepth):
                         path.append(c)
                         yield path
 
diff --git a/test/test_dirtytiles.py b/test/test_dirtytiles.py
index 0002b0c..89ae5c9 100644
--- a/test/test_dirtytiles.py
+++ b/test/test_dirtytiles.py
@@ -59,6 +59,19 @@ class DirtyTilesTest(unittest.TestCase):
         # Make sure they were all returned
         self.assertEqual(len(dirty), 0)
 
+    def test_iterate_levelmax(self):
+        """Same as test_iterate, but specifies the level explicitly"""
+        dirty = set(self.dirty_paths)
+        for p in self.tree.iterate_dirty(3):
+            # Can't use assertIn, was only added in 2.7
+            self.assertTrue(p in dirty)
+
+            # Should not see this one again
+            dirty.remove(p)
+
+        # Make sure they were all returned
+        self.assertEqual(len(dirty), 0)
+
     def test_iterate_fail(self):
         """Meta-test: Make sure test_iterate() would actually fail"""
         # if an extra item were returned"""
@@ -114,7 +127,7 @@ class DirtyTilesTest(unittest.TestCase):
         for p in self.tree.iterate_dirty(2):
             self.assertTrue(p in l2, "%s was not supposed to be returned!" % (p,))
             l2.remove(p)
-        self.assertEqual(len(dirty), 0, "Never iterated over these items: %s" % l2)
+        self.assertEqual(len(l2), 0, "Never iterated over these items: %s" % l2)
 
         # level 1
         l1 = set()
@@ -123,7 +136,7 @@ class DirtyTilesTest(unittest.TestCase):
         for p in self.tree.iterate_dirty(1):
             self.assertTrue(p in l1, "%s was not supposed to be returned!" % (p,))
             l1.remove(p)
-        self.assertEqual(len(dirty), 0, "Never iterated over these items: %s" % l1)
+        self.assertEqual(len(l1), 0, "Never iterated over these items: %s" % l1)
 
 if __name__ == "__main__":
     unittest.main()