Fixed rownumber for cursor subclasses during iterations

Regression introduced to fix ticket #80. Don't use fetchmany to get the
chunks of values. I did it that way because I was ending up into infinite
recursion calling __iter__ from __iter__: the solution has been the
"while 1: yield next()" idiom.
This commit is contained in:
Daniele Varrazzo 2012-02-23 22:55:13 +00:00
parent 7221ea9ec5
commit ebec522a07
2 changed files with 34 additions and 14 deletions

View File

@ -88,25 +88,17 @@ class DictCursorBase(_cursor):
def __iter__(self):
if self._prefetch:
res = _cursor.fetchmany(self, self.itersize)
if not res:
return
res = _cursor.__iter__(self)
first = res.next()
if self._query_executed:
self._build_index()
if not self._prefetch:
res = _cursor.fetchmany(self, self.itersize)
res = _cursor.__iter__(self)
first = res.next()
for r in res:
yield r
# the above was the first itersize record. the following are
# in a repeated loop.
yield first
while 1:
res = _cursor.fetchmany(self, self.itersize)
if not res:
return
for r in res:
yield r
yield res.next()
class DictConnection(_connection):

View File

@ -60,6 +60,11 @@ class ExtrasDictCursorTests(unittest.TestCase):
self.failUnless(row['foo'] == 'qux')
self.failUnless(row[0] == 'qux')
@skip_before_postgres(8, 0)
def testDictCursorWithPlainCursorIterRowNumber(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
self._testIterRowNumber(curs)
def _testWithPlainCursor(self, getter):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
curs.execute("SELECT * FROM ExtrasDictCursorTests")
@ -87,6 +92,11 @@ class ExtrasDictCursorTests(unittest.TestCase):
return row
self._testWithPlainCursorReal(getter)
@skip_before_postgres(8, 0)
def testDictCursorWithPlainCursorRealIterRowNumber(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
self._testIterRowNumber(curs)
def _testWithPlainCursorReal(self, getter):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
curs.execute("SELECT * FROM ExtrasDictCursorTests")
@ -117,6 +127,11 @@ class ExtrasDictCursorTests(unittest.TestCase):
curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.DictCursor)
self._testNamedCursorNotGreedy(curs)
@skip_before_postgres(8, 0)
def testDictCursorWithNamedCursorIterRowNumber(self):
curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.DictCursor)
self._testIterRowNumber(curs)
def _testWithNamedCursor(self, getter):
curs = self.conn.cursor('aname', cursor_factory=psycopg2.extras.DictCursor)
curs.execute("SELECT * FROM ExtrasDictCursorTests")
@ -148,6 +163,11 @@ class ExtrasDictCursorTests(unittest.TestCase):
curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.RealDictCursor)
self._testNamedCursorNotGreedy(curs)
@skip_before_postgres(8, 0)
def testDictCursorRealWithNamedCursorIterRowNumber(self):
curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.RealDictCursor)
self._testIterRowNumber(curs)
def _testWithNamedCursorReal(self, getter):
curs = self.conn.cursor('aname', cursor_factory=psycopg2.extras.RealDictCursor)
curs.execute("SELECT * FROM ExtrasDictCursorTests")
@ -167,6 +187,14 @@ class ExtrasDictCursorTests(unittest.TestCase):
self.assert_(recs[1]['ts'] - recs[0]['ts'] < timedelta(seconds=0.005))
self.assert_(recs[2]['ts'] - recs[1]['ts'] > timedelta(seconds=0.0099))
def _testIterRowNumber(self, curs):
# Only checking for dataset < itersize:
# see CursorTests.test_iter_named_cursor_rownumber
curs.itersize = 20
curs.execute("""select * from generate_series(1,10)""")
for i, r in enumerate(curs):
self.assertEqual(i + 1, curs.rownumber)
class NamedTupleCursorTest(unittest.TestCase):
def setUp(self):