Simplify write in ChunkOutput

Author: Zhen

neo4j/bolt/io.py

@@ -87,12 +87,8 @@ def read(self, n):
             else:
                 value = bytearray(self._view[start:end])
             n -= end - start
-            continue
         return memoryview(value)
 
-    def close(self):
-        self._view = None
-
 
 class ChunkedInputBuffer(object):
 
@@ -213,7 +209,6 @@ def discard_message(self):
             self._frame.close()
             self._origin = self._limit
             self._limit = -1
-            self._frame.close()  # close the frame to release the reference
             self._frame = None
 
 
@@ -236,25 +231,23 @@ def clear(self):
         self._data[0:2] = b"\x00\x00"
 
     def write(self, b):
-        new_data_start = 0
-        new_data_size = len(b)
-        data_size = self._end - self._start
-        if data_size > new_data_size:
-            new_end = self._end + new_data_size
-            self._data[self._end:new_end] = bytearray(data_size)
-        while new_data_start < new_data_size:
-            chunk_occupied = self._end - self._start
-            chunk_remaining = self._max_chunk_size - chunk_occupied
-            if chunk_remaining == 0:
+        to_write = len(b)
+        max_chunk_size = self._max_chunk_size
+        pos = 0
+        while to_write > 0:
+            chunk_size = self._end - self._start
+            remaining = max_chunk_size - chunk_size
+            if remaining == 0 or remaining < to_write <= max_chunk_size:
                 self.chunk()
-                chunk_remaining = self._max_chunk_size
-            chunk_write_size = min(chunk_remaining, new_data_size - new_data_start)
-            new_end = self._end + chunk_write_size
-            new_chunk_size = new_end - self._start
-            self._data[self._end:new_end] = b[new_data_start:(new_data_start + chunk_write_size)]
-            new_data_start += chunk_write_size
-            self._end = new_end
-            self._data[self._header:(self._header + 2)] = struct_pack(">H", new_chunk_size)
+            else:
+                wrote = min(to_write, remaining)
+                new_end = self._end + wrote
+                self._data[self._end:new_end] = b[pos:pos+wrote]
+                self._end = new_end
+                pos += wrote
+                new_chunk_size = self._end - self._start
+                self._data[self._header:(self._header + 2)] = struct_pack(">H", new_chunk_size)
+                to_write -= wrote
 
     def chunk(self):
         self._header = self._end

test/integration/test_session.py

@@ -23,7 +23,7 @@
 from neo4j.v1 import \
     READ_ACCESS, WRITE_ACCESS, \
     CypherError, SessionError, TransactionError, \
-    Node, Relationship, Path, GraphDatabase
+    Node, Relationship, Path
 from neo4j.exceptions import CypherSyntaxError
 
 from test.integration.tools import DirectIntegrationTestCase
@@ -459,12 +459,10 @@ def test_errors_on_begin_transaction(self):
             session.begin_transaction()
 
     def test_large_values(self):
-        driver = GraphDatabase.driver(self.bolt_uri, auth=self.auth_token)
         for i in range(1, 7):
-            session = driver.session()
+            session = self.driver.session()
             session.run("RETURN '{}'".format("A" * 2 ** 20))
             session.close()
-        driver.close()
 
 class TransactionCommittedTestCase(DirectIntegrationTestCase):