Run with unmodified Ariadne

Author: patrys

mysite/routing.py

@@ -1,13 +1,14 @@
 import asyncio
 
-from ariadne import gql, ResolverMap, Subscription
+from ariadne import gql, ResolverMap
 from ariadne.executable_schema import make_executable_schema
 from channels.routing import ProtocolTypeRouter, URLRouter
 from django.conf.urls import url
 from graphql.pyutils import EventEmitter, EventEmitterAsyncIterator
 from graphql.subscription import subscribe
 
 from .graphql import GraphQLHTTPConsumer, GraphQLWebsocketConsumer
+from .subscription import SubscriptionAwareResolverMap
 
 SCHEMA = gql(
     """
@@ -27,7 +28,7 @@
 mutation = ResolverMap("Mutation")
 pubsub = EventEmitter()
 query = ResolverMap("Query")
-messages = Subscription("messages")
+subscription = SubscriptionAwareResolverMap("Subscription")
 
 
 @query.field("hello")
@@ -42,17 +43,18 @@ async def send_message(root, info, message):
     return True
 
 
-@messages.subscriber
+@subscription.subscription("messages")
 def subscribe_messages(root, info):
     return EventEmitterAsyncIterator(pubsub, "message")
 
 
-@messages.resolver
+@subscription.field("messages")
 def push_message(message, info):
     return message
 
 
-schema = make_executable_schema(SCHEMA, [messages, mutation, query])
+schema = make_executable_schema(SCHEMA, [mutation, query, subscription])
+
 
 application = ProtocolTypeRouter(
     {

mysite/subscription.py

@@ -0,0 +1,51 @@
+from typing import Any, AsyncIterator, Callable, Dict, overload
+
+from ariadne import ResolverMap
+from graphql import GraphQLSchema
+
+Subscriber = Callable[..., AsyncIterator]
+
+
+class SubscriptionAwareResolverMap(ResolverMap):
+    _subscribers: Dict[str, Subscriber]
+
+    def __init__(self, name: str):
+        super().__init__(name)
+        self._subscribers = {}
+
+    @overload
+    def subscription(self, name: str) -> Callable[[Subscriber], Subscriber]:
+        pass  # pragma: no cover
+
+    @overload
+    def subscription(  # pylint: disable=function-redefined
+        self, name: str, *, subscriber: Subscriber
+    ) -> Subscriber:  # pylint: disable=function-redefined
+        pass  # pragma: no cover
+
+    def subscription(
+        self, name, *, subscriber=None
+    ):  # pylint: disable=function-redefined
+        if not subscriber:
+            return self.create_register_subscriber(name)
+        self._subscribers[name] = subscriber
+        return subscriber
+
+    def create_register_subscriber(
+        self, name: str
+    ) -> Callable[[Subscriber], Subscriber]:
+        def register_subscriber(f: Subscriber) -> Subscriber:
+            self._subscribers[name] = f
+            return f
+
+        return register_subscriber
+
+    def bind_to_schema(self, schema: GraphQLSchema) -> None:
+        super().bind_to_schema(schema)
+        graphql_type = schema.type_map.get(self.name)
+        for field, subscriber in self._subscribers.items():
+            if field not in graphql_type.fields:
+                raise ValueError(
+                    "Field %s is not defined on type %s" % (field, self.name)
+                )
+            graphql_type.fields[field].subscribe = subscriber