Put ListenerQueue in separate file

Author: leszekhanusz

gql/transport/aiohttp_websockets.py

@@ -22,72 +22,19 @@
 from graphql import DocumentNode, ExecutionResult, print_ast
 from multidict import CIMultiDictProxy
 
-from gql.transport.aiohttp import AIOHTTPTransport
-from gql.transport.async_transport import AsyncTransport
-from gql.transport.exceptions import (
+from .aiohttp import AIOHTTPTransport
+from .async_transport import AsyncTransport
+from .exceptions import (
     TransportAlreadyConnected,
     TransportClosed,
     TransportProtocolError,
     TransportQueryError,
     TransportServerError,
 )
+from .websockets_common import ListenerQueue
 
 log = logging.getLogger("gql.transport.aiohttp_websockets")
 
-ParsedAnswer = Tuple[str, Optional[ExecutionResult]]
-
-
-class ListenerQueue:
-    """Special queue used for each query waiting for server answers
-
-    If the server is stopped while the listener is still waiting,
-    Then we send an exception to the queue and this exception will be raised
-    to the consumer once all the previous messages have been consumed from the queue
-    """
-
-    def __init__(self, query_id: int, send_stop: bool) -> None:
-        self.query_id: int = query_id
-        self.send_stop: bool = send_stop
-        self._queue: asyncio.Queue = asyncio.Queue()
-        self._closed: bool = False
-
-    async def get(self) -> ParsedAnswer:
-
-        try:
-            item = self._queue.get_nowait()
-        except asyncio.QueueEmpty:
-            item = await self._queue.get()
-
-        self._queue.task_done()
-
-        # If we receive an exception when reading the queue, we raise it
-        if isinstance(item, Exception):
-            self._closed = True
-            raise item
-
-        # Don't need to save new answers or
-        # send the stop message if we already received the complete message
-        answer_type, execution_result = item
-        if answer_type == "complete":
-            self.send_stop = False
-            self._closed = True
-
-        return item
-
-    async def put(self, item: ParsedAnswer) -> None:
-
-        if not self._closed:
-            await self._queue.put(item)
-
-    async def set_exception(self, exception: Exception) -> None:
-
-        # Put the exception in the queue
-        await self._queue.put(exception)
-
-        # Don't need to send stop messages in case of error
-        self.send_stop = False
-        self._closed = True
-
 
 class AIOHTTPWebsocketsTransport(AsyncTransport):
 

gql/transport/websockets_base.py

@@ -21,63 +21,10 @@
     TransportQueryError,
     TransportServerError,
 )
+from .websockets_common import ListenerQueue
 
 log = logging.getLogger("gql.transport.websockets")
 
-ParsedAnswer = Tuple[str, Optional[ExecutionResult]]
-
-
-class ListenerQueue:
-    """Special queue used for each query waiting for server answers
-
-    If the server is stopped while the listener is still waiting,
-    Then we send an exception to the queue and this exception will be raised
-    to the consumer once all the previous messages have been consumed from the queue
-    """
-
-    def __init__(self, query_id: int, send_stop: bool) -> None:
-        self.query_id: int = query_id
-        self.send_stop: bool = send_stop
-        self._queue: asyncio.Queue = asyncio.Queue()
-        self._closed: bool = False
-
-    async def get(self) -> ParsedAnswer:
-
-        try:
-            item = self._queue.get_nowait()
-        except asyncio.QueueEmpty:
-            item = await self._queue.get()
-
-        self._queue.task_done()
-
-        # If we receive an exception when reading the queue, we raise it
-        if isinstance(item, Exception):
-            self._closed = True
-            raise item
-
-        # Don't need to save new answers or
-        # send the stop message if we already received the complete message
-        answer_type, execution_result = item
-        if answer_type == "complete":
-            self.send_stop = False
-            self._closed = True
-
-        return item
-
-    async def put(self, item: ParsedAnswer) -> None:
-
-        if not self._closed:
-            await self._queue.put(item)
-
-    async def set_exception(self, exception: Exception) -> None:
-
-        # Put the exception in the queue
-        await self._queue.put(exception)
-
-        # Don't need to send stop messages in case of error
-        self.send_stop = False
-        self._closed = True
-
 
 class WebsocketsTransportBase(AsyncTransport):
     """abstract :ref:`Async Transport <async_transports>` used to implement

gql/transport/websockets_common/__init__.py

@@ -0,0 +1,3 @@
+from .listener_queue import ListenerQueue, ParsedAnswer
+
+__all__ = ["ListenerQueue", "ParsedAnswer"]

gql/transport/websockets_common/listener_queue.py

@@ -0,0 +1,58 @@
+import asyncio
+from typing import Optional, Tuple
+
+from graphql import ExecutionResult
+
+ParsedAnswer = Tuple[str, Optional[ExecutionResult]]
+
+
+class ListenerQueue:
+    """Special queue used for each query waiting for server answers
+
+    If the server is stopped while the listener is still waiting,
+    Then we send an exception to the queue and this exception will be raised
+    to the consumer once all the previous messages have been consumed from the queue
+    """
+
+    def __init__(self, query_id: int, send_stop: bool) -> None:
+        self.query_id: int = query_id
+        self.send_stop: bool = send_stop
+        self._queue: asyncio.Queue = asyncio.Queue()
+        self._closed: bool = False
+
+    async def get(self) -> ParsedAnswer:
+
+        try:
+            item = self._queue.get_nowait()
+        except asyncio.QueueEmpty:
+            item = await self._queue.get()
+
+        self._queue.task_done()
+
+        # If we receive an exception when reading the queue, we raise it
+        if isinstance(item, Exception):
+            self._closed = True
+            raise item
+
+        # Don't need to save new answers or
+        # send the stop message if we already received the complete message
+        answer_type, execution_result = item
+        if answer_type == "complete":
+            self.send_stop = False
+            self._closed = True
+
+        return item
+
+    async def put(self, item: ParsedAnswer) -> None:
+
+        if not self._closed:
+            await self._queue.put(item)
+
+    async def set_exception(self, exception: Exception) -> None:
+
+        # Put the exception in the queue
+        await self._queue.put(exception)
+
+        # Don't need to send stop messages in case of error
+        self.send_stop = False
+        self._closed = True