ResultIter can now deal with generators of resultproxies

dataset/persistence/util.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from inspect import isgenerator
 
 from sqlalchemy import Integer, UnicodeText, Float, DateTime, Boolean
 
@@ -16,21 +17,36 @@ def guess_type(sample):
 
 
 class ResultIter(object):
-    """ SQLAlchemy ResultProxies are not iterable to get a 
+    """ SQLAlchemy ResultProxies are not iterable to get a
     list of dictionaries. This is to wrap them. """
 
-    def __init__(self, rp):
-        self.rp = rp
-        self.count = rp.rowcount
-        self.keys = self.rp.keys()
+    def __init__(self, result_proxies):
+        if not isgenerator(result_proxies):
+            result_proxies = iter((result_proxies, ))
+        self.result_proxies = result_proxies
+
+        self.count = 0
+        if not self._next_rp():
+            raise StopIteration
+
+    def _next_rp(self):
+        try:
+            self.rp = self.result_proxies.next()
+            self.count += self.rp.rowcount
+            self.keys = self.rp.keys()
+            return True
+        except StopIteration:
+            return False
 
     def next(self):
         row = self.rp.fetchone()
         if row is None:
-            raise StopIteration
+            if self._next_rp():
+                return self.next()
+            else:
+                # stop here
+                raise StopIteration
         return dict(zip(self.keys, row))
 
     def __iter__(self):
         return self
-
-