From 73ff6bdc2324df660b9131db1a14dd0ec5083437 Mon Sep 17 00:00:00 2001
From: davebshow <davebshow@gmail.com>
Date: Sun, 3 Jul 2016 01:57:30 -0400
Subject: [PATCH] added a force close option for async iterator
---
goblin/gremlin_python_driver/driver.py | 5 ++++-
tests/test_engine.py | 1 -
2 files changed, 4 insertions(+), 2 deletions(-)
diff --git a/goblin/gremlin_python_driver/driver.py b/goblin/gremlin_python_driver/driver.py
index 33a723d..dc8b206 100644
--- a/goblin/gremlin_python_driver/driver.py
+++ b/goblin/gremlin_python_driver/driver.py
@@ -48,10 +48,11 @@ class Driver:
class AsyncResponseIter:
- def __init__(self, ws, loop, conn):
+ def __init__(self, ws, loop, conn, *, force_close=True):
self._ws = ws
self._loop = loop
self._conn = conn
+ self._force_close = force_close
self._closed = False
async def __aiter__(self):
@@ -62,6 +63,8 @@ class AsyncResponseIter:
if msg:
return msg
else:
+ if self._force_close:
+ await self.close()
raise StopAsyncIteration
async def close(self):
diff --git a/tests/test_engine.py b/tests/test_engine.py
index 4081dd6..c9bf716 100644
--- a/tests/test_engine.py
+++ b/tests/test_engine.py
@@ -85,7 +85,6 @@ class TestEngine(unittest.TestCase):
results += msg.data
print(len(results))
self.assertTrue(len(results) > 1)
- await stream.close()
await engine.close()
self.loop.run_until_complete(go())
--
GitLab