Support optional iterator protocol on cursors (#38)

Use the original iterator if the cursor is already iterable,
make it iterable otherwise.

Author: Cito

dbutils/steady_db.py

@@ -542,6 +542,14 @@ def __exit__(self, *exc):
         """Exit the runtime context for the cursor object."""
         self.close()
 
+    def __iter__(self):
+        """Make cursor compatible to the iteration protocol."""
+        cursor = self._cursor
+        try:  # use iterator provided by original cursor
+            return iter(cursor)
+        except TypeError:  # create iterator if not provided
+            return iter(cursor.fetchone, None)
+
     def setinputsizes(self, sizes):
         """Store input sizes in case cursor needs to be reopened."""
         self._inputsizes = sizes

tests/test_steady_db.py

@@ -273,6 +273,36 @@ def test_cursor_context_handler(self):
             self.assertEqual(cursor.fetchone(), 'test')
         self.assertEqual(db._con.open_cursors, 0)
 
+    def test_cursor_as_iterator_provided(self):
+        db = SteadyDBconnect(
+            dbapi, 0, None, None, None, True,
+            'SteadyDBTestDB', user='SteadyDBTestUser')
+        self.assertEqual(db._con.open_cursors, 0)
+        cursor = db.cursor()
+        self.assertEqual(db._con.open_cursors, 1)
+        cursor.execute('select test')
+        _cursor = cursor._cursor
+        try:
+            assert not hasattr(_cursor, 'iter')
+            _cursor.__iter__ = lambda: ['test-iter']
+            assert list(iter(cursor)) == ['test']
+        finally:
+            del _cursor.__iter__
+        cursor.close()
+        self.assertEqual(db._con.open_cursors, 0)
+
+    def test_cursor_as_iterator_created(self):
+        db = SteadyDBconnect(
+            dbapi, 0, None, None, None, True,
+            'SteadyDBTestDB', user='SteadyDBTestUser')
+        self.assertEqual(db._con.open_cursors, 0)
+        cursor = db.cursor()
+        self.assertEqual(db._con.open_cursors, 1)
+        cursor.execute('select test')
+        assert list(iter(cursor)) == ['test']
+        cursor.close()
+        self.assertEqual(db._con.open_cursors, 0)
+
     def test_connection_creator_function(self):
         db1 = SteadyDBconnect(
             dbapi, 0, None, None, None, True,