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.

lib/extras.py

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

tests/test_extras_dictcursor.py

@@ -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):