test_subscribe: Replace EventEmitter with SimplePubSub

Replicates graphql/graphql-js@156c76e

Author: Cito

.coveragerc

@@ -10,6 +10,8 @@ omit =
 [report]
 exclude_lines =
     pragma: no cover
+    except ImportError:
+    \# Python <
     raise NotImplementedError
     raise TypeError\(f?"Unexpected
     assert False,

docs/conf.py

@@ -108,7 +108,7 @@
     'execution': ['execute', 'middleware'],
     'language': ['ast', 'directive_locations', 'location',
                  'source', 'token_kind', 'visitor'],
-    'pyutils': ['event_emitter', 'frozen_list', 'path'],
+    'pyutils': ['simple_pub_sub', 'frozen_list', 'path'],
     'subscription': [],
     'type': ['definition', 'directives', 'schema'],
     'utilities': ['find_breaking_changes', 'type_info'],
@@ -129,7 +129,9 @@
 ignore_references = set('''
 GNT GT T
 enum.Enum
+traceback
 asyncio.events.AbstractEventLoop
+graphql.subscription.map_async_iterator.MapAsyncIterator
 graphql.type.schema.InterfaceImplementations
 graphql.validation.validation_context.VariableUsage
 graphql.validation.rules.known_argument_names.KnownArgumentNamesOnDirectivesRule

docs/modules/pyutils.rst

@@ -13,8 +13,8 @@ PyUtils
 .. autofunction:: did_you_mean
 .. autofunction:: register_description
 .. autofunction:: unregister_description
-.. autoclass:: EventEmitter
-.. autoclass:: EventEmitterAsyncIterator
+.. autoclass:: SimplePubSub
+.. autoclass:: SimplePubSubIterator
 .. autofunction:: identity_func
 .. autofunction:: inspect
 .. autofunction:: is_finite

docs/modules/subscription.rst

@@ -4,4 +4,14 @@ Subscription
 .. currentmodule:: graphql.subscription
 
 .. automodule:: graphql.subscription
+   :no-members:
+   :no-inherited-members:
 
+.. autofunction:: subscribe
+
+Helpers
+-------
+
+.. autofunction:: create_source_event_stream
+
+.. autoclass:: MapAsyncIterator

src/graphql/__init__.py

@@ -274,7 +274,7 @@
     MiddlewareManager,
 )
 
-from .subscription import subscribe, create_source_event_stream
+from .subscription import subscribe, create_source_event_stream, MapAsyncIterator
 
 
 # Validate GraphQL queries.
@@ -622,6 +622,7 @@
     "MiddlewareManager",
     "subscribe",
     "create_source_event_stream",
+    "MapAsyncIterator",
     "validate",
     "ValidationContext",
     "ValidationRule",

src/graphql/pyutils/__init__.py

@@ -17,7 +17,6 @@
     unregister_description,
 )
 from .did_you_mean import did_you_mean
-from .event_emitter import EventEmitter, EventEmitterAsyncIterator
 from .identity_func import identity_func
 from .inspect import inspect
 from .is_awaitable import is_awaitable
@@ -31,6 +30,7 @@
 from .frozen_dict import FrozenDict
 from .path import Path
 from .print_path_list import print_path_list
+from .simple_pub_sub import SimplePubSub, SimplePubSubIterator
 from .undefined import Undefined, UndefinedType
 
 __all__ = [
@@ -42,8 +42,6 @@
     "is_description",
     "register_description",
     "unregister_description",
-    "EventEmitter",
-    "EventEmitterAsyncIterator",
     "identity_func",
     "inspect",
     "is_awaitable",
@@ -57,6 +55,8 @@
     "FrozenDict",
     "Path",
     "print_path_list",
+    "SimplePubSub",
+    "SimplePubSubIterator",
     "Undefined",
     "UndefinedType",
 ]

src/graphql/pyutils/event_emitter.py

@@ -0,0 +1,81 @@
+from asyncio import Future, Queue, ensure_future, sleep
+from inspect import isawaitable
+from typing import Any, AsyncIterator, Callable, Optional, Set
+
+try:
+    from asyncio import get_running_loop
+except ImportError:
+    from asyncio import get_event_loop as get_running_loop  # Python < 3.7
+
+
+__all__ = ["SimplePubSub", "SimplePubSubIterator"]
+
+
+class SimplePubSub:
+    """A very simple publish-subscript system.
+
+    Creates an AsyncIterator from an EventEmitter.
+
+    Useful for mocking a PubSub system for tests.
+    """
+
+    subscribers: Set[Callable]
+
+    def __init__(self) -> None:
+        self.subscribers = set()
+
+    def emit(self, event: Any) -> bool:
+        """Emit an event."""
+        for subscriber in self.subscribers:
+            result = subscriber(event)
+            if isawaitable(result):
+                ensure_future(result)
+        return bool(self.subscribers)
+
+    def get_subscriber(
+        self, transform: Optional[Callable] = None
+    ) -> "SimplePubSubIterator":
+        return SimplePubSubIterator(self, transform)
+
+
+class SimplePubSubIterator(AsyncIterator):
+    def __init__(self, pubsub: SimplePubSub, transform: Optional[Callable]) -> None:
+        self.pubsub = pubsub
+        self.transform = transform
+        self.pull_queue: Queue[Future] = Queue()
+        self.push_queue: Queue[Any] = Queue()
+        self.listening = True
+        pubsub.subscribers.add(self.push_value)
+
+    def __aiter__(self) -> "SimplePubSubIterator":
+        return self
+
+    async def __anext__(self) -> Any:
+        if not self.listening:
+            raise StopAsyncIteration
+        await sleep(0)
+        if not self.push_queue.empty():
+            return await self.push_queue.get()
+        future = get_running_loop().create_future()
+        await self.pull_queue.put(future)
+        return future
+
+    async def aclose(self) -> None:
+        if self.listening:
+            await self.empty_queue()
+
+    async def empty_queue(self) -> None:
+        self.listening = False
+        self.pubsub.subscribers.remove(self.push_value)
+        while not self.pull_queue.empty():
+            future = await self.pull_queue.get()
+            future.cancel()
+        while not self.push_queue.empty():
+            await self.push_queue.get()
+
+    async def push_value(self, event: Any) -> None:
+        value = event if self.transform is None else self.transform(event)
+        if self.pull_queue.empty():
+            await self.push_queue.put(value)
+        else:
+            (await self.pull_queue.get()).set_result(value)

src/graphql/pyutils/simple_pub_sub.py

@@ -5,5 +5,6 @@
 """
 
 from .subscribe import subscribe, create_source_event_stream
+from .map_async_iterator import MapAsyncIterator
 
-__all__ = ["subscribe", "create_source_event_stream"]
+__all__ = ["subscribe", "create_source_event_stream", "MapAsyncIterator"]

src/graphql/subscription/__init__.py

@@ -30,9 +30,11 @@ def __init__(
         self._close_event = Event()
 
     def __aiter__(self) -> "MapAsyncIterator":
+        """Get the iterator object."""
         return self
 
     async def __anext__(self) -> Any:
+        """Get the next value of the iterator."""
         if self.is_closed:
             if not isasyncgen(self.iterator):
                 raise StopAsyncIteration
@@ -79,6 +81,7 @@ async def athrow(
         value: Optional[BaseException] = None,
         traceback: Optional[TracebackType] = None,
     ) -> None:
+        """Throw an exception into the asynchronous iterator."""
         if not self.is_closed:
             athrow = getattr(self.iterator, "athrow", None)
             if athrow:
@@ -98,6 +101,7 @@ async def athrow(
                 raise value
 
     async def aclose(self) -> None:
+        """Close the iterator."""
         if not self.is_closed:
             aclose = getattr(self.iterator, "aclose", None)
             if aclose:
@@ -109,10 +113,12 @@ async def aclose(self) -> None:
 
     @property
     def is_closed(self) -> bool:
+        """Check whether the iterator is closed."""
         return self._close_event.is_set()
 
     @is_closed.setter
     def is_closed(self, value: bool) -> None:
+        """Mark the iterator as closed."""
         if value:
             self._close_event.set()
         else: