From 187ed2680f3f037bdc691357fb3c2f0bcff02143 Mon Sep 17 00:00:00 2001 From: Gregor Aisch Date: Fri, 12 Apr 2013 16:41:17 +0200 Subject: [PATCH] ResultIter can now deal with generators of resultproxies --- dataset/persistence/util.py | 32 ++++++++++++++++++++++++-------- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/dataset/persistence/util.py b/dataset/persistence/util.py index bd902c1..019f620 100644 --- a/dataset/persistence/util.py +++ b/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 - -