[pypy-commit] pypy default: test and fix for sqlite cursor iteration, simplify/optimize

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: 
Changeset: r62184:7e1f21fd7d13
Date: 2013-03-07 15:27 -0500
http://bitbucket.org/pypy/pypy/changeset/7e1f21fd7d13/

Log:	test and fix for sqlite cursor iteration, simplify/optimize

diff --git a/lib_pypy/_sqlite3.py b/lib_pypy/_sqlite3.py
--- a/lib_pypy/_sqlite3.py
+++ b/lib_pypy/_sqlite3.py
@@ -919,24 +919,24 @@
                     "Cursor needed to be reset because of commit/rollback "
                     "and can no longer be fetched from.")
 
-    # do all statements
-    def fetchone(self):
+    def __iter__(self):
+        return self
+
+    def __next__(self):
         self.__check_cursor()
         self.__check_reset()
+        if not self.__statement:
+            raise StopIteration
+        return self.__statement._next(self)
 
-        if self.__statement is None:
-            return None
+    if sys.version_info[0] < 3:
+        next = __next__
+        del __next__
 
-        try:
-            return self.__statement._next(self)
-        except StopIteration:
-            return None
+    def fetchone(self):
+        return next(self, None)
 
     def fetchmany(self, size=None):
-        self.__check_cursor()
-        self.__check_reset()
-        if self.__statement is None:
-            return []
         if size is None:
             size = self.arraysize
         lst = []
@@ -947,15 +947,8 @@
         return lst
 
     def fetchall(self):
-        self.__check_cursor()
-        self.__check_reset()
-        if self.__statement is None:
-            return []
         return list(self)
 
-    def __iter__(self):
-        return iter(self.fetchone, None)
-
     def __get_connection(self):
         return self.__connection
     connection = property(__get_connection)
diff --git a/pypy/module/test_lib_pypy/test_sqlite3.py b/pypy/module/test_lib_pypy/test_sqlite3.py
--- a/pypy/module/test_lib_pypy/test_sqlite3.py
+++ b/pypy/module/test_lib_pypy/test_sqlite3.py
@@ -52,6 +52,16 @@
     # raises ProgrammingError because should check closed before check args
     pytest.raises(_sqlite3.ProgrammingError, "con()")
 
+def test_cursor_iter():
+    con = _sqlite3.connect(':memory:')
+    cur = con.cursor()
+    with pytest.raises(StopIteration):
+        next(cur)
+    cur = con.execute('select 1')
+    next(cur)
+    with pytest.raises(StopIteration):
+        next(cur)
+
 def test_cursor_after_close():
      con = _sqlite3.connect(':memory:')
      cur = con.execute('select 1')