Iteration fixes

Author: technige

examples/test_examples.py

@@ -158,7 +158,7 @@ def test_transaction_commit(self):
         tx.commit()
         # end::transaction-commit[]
         result = session.run("MATCH (p:Person {name: 'Guinevere'}) RETURN count(p)")
-        record = next(result)
+        record = next(iter(result))
         assert record["count(p)"] == 1
         session.close()
 
@@ -171,7 +171,7 @@ def test_transaction_rollback(self):
         tx.rollback()
         # end::transaction-rollback[]
         result = session.run("MATCH (p:Person {name: 'Merlin'}) RETURN count(p)")
-        record = next(result)
+        record = next(iter(result))
         assert record["count(p)"] == 0
         session.close()
 

neo4j/v1/session.py

@@ -206,19 +206,14 @@ def on_failure(metadata):
         pull_all_response.on_failure = on_failure
 
     def __iter__(self):
-        return self
-
-    def __next__(self):
-        if self._buffer:
+        while self._buffer and not self._consumed:
             values = self._buffer.popleft()
-            return Record(self.keys(), tuple(map(hydrated, values)))
-        elif self._consumed:
-            raise StopIteration()
-        else:
-            fetch = self.connection.fetch
-            while not self._buffer and not self._consumed:
-                fetch()
-            return self.__next__()
+            yield Record(self.keys(), tuple(map(hydrated, values)))
+        while not self._consumed:
+            self.connection.fetch()
+            if self._buffer:
+                values = self._buffer.popleft()
+                yield Record(self.keys(), tuple(map(hydrated, values)))
 
     def keys(self):
         """ Return the keys for the records.

test/tck/tck_util.py

@@ -29,15 +29,17 @@
 def send_string(text):
     session = driver.session()
     result = session.run(text)
+    records = list(result)
     session.close()
-    return list(result)
+    return records
 
 
 def send_parameters(statement, parameters):
     session = driver.session()
     result = session.run(statement, parameters)
+    records = list(result)
     session.close()
-    return list(result)
+    return records
 
 
 try:

test/test_session.py

@@ -330,7 +330,7 @@ def test_automatic_reset_after_failure(self):
                 session.run("X").consume()
             except CypherError:
                 result = session.run("RETURN 1")
-                record = next(result)
+                record = next(iter(result))
                 assert record[0] == 1
             else:
                 assert False, "A Cypher error should have occurred"
@@ -424,7 +424,7 @@ def test_can_commit_transaction(self):
 
             # Create a node
             result = tx.run("CREATE (a) RETURN id(a)")
-            record = next(result)
+            record = next(iter(result))
             node_id = record[0]
             assert isinstance(node_id, int)
 
@@ -437,7 +437,7 @@ def test_can_commit_transaction(self):
             # Check the property value
             result = session.run("MATCH (a) WHERE id(a) = {n} "
                                  "RETURN a.foo", {"n": node_id})
-            record = next(result)
+            record = next(iter(result))
             value = record[0]
             assert value == "bar"
 
@@ -447,7 +447,7 @@ def test_can_rollback_transaction(self):
 
             # Create a node
             result = tx.run("CREATE (a) RETURN id(a)")
-            record = next(result)
+            record = next(iter(result))
             node_id = record[0]
             assert isinstance(node_id, int)
 
@@ -467,7 +467,7 @@ def test_can_commit_transaction_using_with_block(self):
             with session.begin_transaction() as tx:
                 # Create a node
                 result = tx.run("CREATE (a) RETURN id(a)")
-                record = next(result)
+                record = next(iter(result))
                 node_id = record[0]
                 assert isinstance(node_id, int)
 
@@ -480,7 +480,7 @@ def test_can_commit_transaction_using_with_block(self):
             # Check the property value
             result = session.run("MATCH (a) WHERE id(a) = {n} "
                                  "RETURN a.foo", {"n": node_id})
-            record = next(result)
+            record = next(iter(result))
             value = record[0]
             assert value == "bar"
 
@@ -489,7 +489,7 @@ def test_can_rollback_transaction_using_with_block(self):
             with session.begin_transaction() as tx:
                 # Create a node
                 result = tx.run("CREATE (a) RETURN id(a)")
-                record = next(result)
+                record = next(iter(result))
                 node_id = record[0]
                 assert isinstance(node_id, int)