Added a flat page iterator for arbitrary page lists.

[automanga.git] / manga / lib.py

diff --git a/manga/lib.py b/manga/lib.py
index 712de71..7988438 100644 (file)
--- a/manga/lib.py
+++ b/manga/lib.py
@@ -9,3 +9,31 @@ class manga(pagelist):
 
 class page(object):
     pass
+
+class pageiter(object):
+    def __init__(self, root):
+        self.nstack = [0]
+        self.lstack = [root]
+
+    def next(self):
+        while True:
+            if len(self.nstack) == 0:
+                raise StopIteration
+            try:
+                node = self.lstack[-1][self.nstack[-1]]
+            except IndexError:
+                self.lstack.pop()
+                self.nstack.pop()
+                if len(self.nstack) > 0:
+                    self.nstack[-1] += 1
+                continue
+            if isinstance(node, page):
+                nl = tuple(self.nstack)
+                self.nstack[-1] += 1
+                return nl, node
+            elif isinstance(node, pagelist):
+                self.lstack.append(node)
+                self.nstack.append(0)
+
+    def __iter__(self):
+        return self