Commits

Brian Kearns committed 7e1f21f

test and fix for sqlite cursor iteration, simplify/optimize

  • Participants
  • Parent commits 8ce8657

Comments (0)

Files changed (2)

lib_pypy/_sqlite3.py

                     "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 = []
         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)

pypy/module/test_lib_pypy/test_sqlite3.py

     # 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')