Fix type_adapter for subscriptions and async transports

Author: leszekhanusz

gql/client.py

@@ -67,16 +67,21 @@ def __init__(
         # Enforced timeout of the execute function
         self.execute_timeout = execute_timeout
 
+        # Fetch schema from transport directly if we are using a sync transport
+        if isinstance(transport, Transport) and fetch_schema_from_transport:
+            with self as session:
+                session.fetch_schema()
+
         # Dictionary where the name of the custom scalar type is the key and the
         # value is a class which has a `parse_value()` function
+        self.custom_types = custom_types
+
+        # Create a type_adapter instance directly here if we received the schema
+        # locally or from a sync transport
         self.type_adapter = (
             TypeAdapter(schema, custom_types) if custom_types and schema else None
         )
 
-        if isinstance(transport, Transport) and fetch_schema_from_transport:
-            with self as session:
-                session.fetch_schema()
-
     def validate(self, document):
         assert (
             self.schema
@@ -211,9 +216,7 @@ def _execute(self, document: DocumentNode, *args, **kwargs) -> ExecutionResult:
 
         return self.transport.execute(document, *args, **kwargs)
 
-    def execute(
-        self, document: DocumentNode, *args, **kwargs
-    ) -> Optional[Dict[str, Any]]:
+    def execute(self, document: DocumentNode, *args, **kwargs) -> Dict[str, Any]:
 
         # Validate and execute on the transport
         result = self._execute(document, *args, **kwargs)
@@ -227,11 +230,9 @@ def execute(
         ), "Transport returned an ExecutionResult without data or errors"
 
         if self.client.type_adapter:
-            result = result._replace(
-                data=self.client.type_adapter.convert_scalars(result.data)
-            )
-
-        return result.data
+            return self.client.type_adapter.convert_scalars(result.data)
+        else:
+            return result.data
 
     def fetch_schema(self) -> None:
         execution_result = self.transport.execute(parse(get_introspection_query()))
@@ -263,6 +264,13 @@ async def fetch_and_validate(self, document: DocumentNode):
         if self.client.fetch_schema_from_transport and not self.client.schema:
             await self.fetch_schema()
 
+            # Once we have received the schema from the async transport,
+            # we can create a TypeAdapter instance if the user provided custom types
+            if self.client.custom_types and self.client.schema:
+                self.client.type_adapter = TypeAdapter(
+                    self.client.schema, self.client.custom_types
+                )
+
         # Validate document
         if self.client.schema:
             self.client.validate(document)
@@ -293,7 +301,7 @@ async def _subscribe(
 
     async def subscribe(
         self, document: DocumentNode, *args, **kwargs
-    ) -> AsyncGenerator[Optional[Dict[str, Any]], None]:
+    ) -> AsyncGenerator[Dict[str, Any], None]:
 
         # Validate and subscribe on the transport
         async for result in self._subscribe(document, *args, **kwargs):
@@ -304,10 +312,9 @@ async def subscribe(
 
             elif result.data is not None:
                 if self.client.type_adapter:
-                    result = result._replace(
-                        data=self.client.type_adapter.convert_scalars(result.data)
-                    )
-                yield result.data
+                    yield self.client.type_adapter.convert_scalars(result.data)
+                else:
+                    yield result.data
 
     async def _execute(
         self, document: DocumentNode, *args, **kwargs
@@ -322,9 +329,7 @@ async def _execute(
             self.client.execute_timeout,
         )
 
-    async def execute(
-        self, document: DocumentNode, *args, **kwargs
-    ) -> Optional[Dict[str, Any]]:
+    async def execute(self, document: DocumentNode, *args, **kwargs) -> Dict[str, Any]:
 
         # Validate and execute on the transport
         result = await self._execute(document, *args, **kwargs)
@@ -338,11 +343,9 @@ async def execute(
         ), "Transport returned an ExecutionResult without data or errors"
 
         if self.client.type_adapter:
-            result = result._replace(
-                data=self.client.type_adapter.convert_scalars(result.data)
-            )
-
-        return result.data
+            return self.client.type_adapter.convert_scalars(result.data)
+        else:
+            return result.data
 
     async def fetch_schema(self) -> None:
         execution_result = await self.transport.execute(

gql/type_adapter.py

@@ -80,6 +80,11 @@ def traverse_schema(
             schema_root = self.schema.query_type
         elif self.schema.mutation_type and keys[0] in self.schema.mutation_type.fields:
             schema_root = self.schema.mutation_type
+        elif (
+            self.schema.subscription_type
+            and keys[0] in self.schema.subscription_type.fields
+        ):
+            schema_root = self.schema.subscription_type
         else:
             return None
 
@@ -95,6 +100,7 @@ def _get_decoded_scalar_type(self, keys: List[str], value: str) -> str:
 
         If it is a custom scalar, return the deserialized value, as
         output by `<CustomScalarType>.parse_value()`"""
+
         scalar_type = self._lookup_scalar_type(keys)
         if scalar_type and scalar_type in self.custom_types:
             return self.custom_types[scalar_type].parse_value(value)

tests/test_async_client_validation.py

@@ -267,3 +267,47 @@ async def test_async_client_validation_fetch_schema_from_server_with_client_argu
 
         with pytest.raises(graphql.error.GraphQLError):
             await session.execute(query)
+
+
+class ToLowercase:
+    @staticmethod
+    def parse_value(value: str):
+        return value.lower()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("server", [hero_server_answers], indirect=True)
+async def test_async_client_validation_fetch_schema_from_server_with_custom_types(
+    event_loop, server
+):
+
+    url = f"ws://{server.hostname}:{server.port}/graphql"
+
+    sample_transport = WebsocketsTransport(url=url)
+
+    custom_types = {"String": ToLowercase}
+
+    async with Client(
+        transport=sample_transport,
+        fetch_schema_from_transport=True,
+        custom_types=custom_types,
+    ) as session:
+
+        query = gql(
+            """
+            query HeroNameQuery {
+              hero {
+                name
+              }
+            }
+        """
+        )
+
+        result = await session.execute(query)
+
+        print("Client received:", result)
+
+        # The expected hero name is now in lowercase
+        expected = {"hero": {"name": "r2-d2"}}
+
+        assert result == expected

tests/test_requests.py

@@ -13,6 +13,8 @@
 )
 from gql.transport.requests import RequestsHTTPTransport
 
+from .test_type_adapter import Capitalize
+
 query1_str = """
     query getContinents {
       continents {
@@ -201,3 +203,57 @@ def test_code():
             sample_transport.execute(query)
 
     await run_sync_test(event_loop, server, test_code)
+
+
+partial_schema = """
+
+type Continent {
+  code: ID!
+  name: String!
+}
+
+type Query {
+  continents: [Continent!]!
+}
+
+"""
+
+
+@pytest.mark.asyncio
+async def test_requests_query_with_custom_types(event_loop, aiohttp_server):
+    async def handler(request):
+        return web.Response(text=query1_server_answer, content_type="application/json")
+
+    app = web.Application()
+    app.router.add_route("POST", "/", handler)
+    server = await aiohttp_server(app)
+
+    url = server.make_url("/")
+
+    def test_code():
+        sample_transport = RequestsHTTPTransport(url=url)
+
+        custom_types = {"String": Capitalize}
+
+        # Instanciate a client which will capitalize all the String scalars
+        with Client(
+            transport=sample_transport,
+            type_def=partial_schema,
+            custom_types=custom_types,
+        ) as session:
+
+            query = gql(query1_str)
+
+            # Execute query synchronously
+            result = session.execute(query)
+
+            continents = result["continents"]
+
+            africa = continents[0]
+
+            assert africa["code"] == "AF"
+
+            # Check that the string is capitalized
+            assert africa["name"] == "AFRICA"
+
+    await run_sync_test(event_loop, server, test_code)

tests/test_websocket_subscription.py

@@ -10,9 +10,24 @@
 
 from .conftest import MS, WebSocketServer
 
-countdown_server_answer = (
-    '{{"type":"data","id":"{query_id}","payload":{{"data":{{"number":{number}}}}}}}'
-)
+countdown_schema = """
+
+type Number {
+  number: Int!
+}
+
+type Query {
+  currentCount: Number!
+}
+
+type Subscription {
+  countdown(count: Int!): Number!
+}
+
+"""
+
+countdown_server_answer = '{{"type":"data","id":"{query_id}","payload":\
+{{"data":{{"countdown":{{"number":{number}}}}}}}}}'
 
 WITH_KEEPALIVE = False
 
@@ -111,7 +126,7 @@ async def test_websocket_subscription(event_loop, client_and_server, subscriptio
 
     async for result in session.subscribe(subscription):
 
-        number = result["number"]
+        number = result["countdown"]["number"]
         print(f"Number received: {number}")
 
         assert number == count
@@ -134,7 +149,7 @@ async def test_websocket_subscription_break(
 
     async for result in session.subscribe(subscription):
 
-        number = result["number"]
+        number = result["countdown"]["number"]
         print(f"Number received: {number}")
 
         assert number == count
@@ -165,7 +180,7 @@ async def task_coro():
         nonlocal count
         async for result in session.subscribe(subscription):
 
-            number = result["number"]
+            number = result["countdown"]["number"]
             print(f"Number received: {number}")
 
             assert number == count
@@ -204,7 +219,7 @@ async def task_coro():
         nonlocal count
         async for result in session.subscribe(subscription):
 
-            number = result["number"]
+            number = result["countdown"]["number"]
             print(f"Number received: {number}")
 
             assert number == count
@@ -269,7 +284,7 @@ async def test_websocket_subscription_server_connection_closed(
 
         async for result in session.subscribe(subscription):
 
-            number = result["number"]
+            number = result["countdown"]["number"]
             print(f"Number received: {number}")
 
             assert number == count
@@ -292,7 +307,7 @@ async def test_websocket_subscription_slow_consumer(
     async for result in session.subscribe(subscription):
         await asyncio.sleep(10 * MS)
 
-        number = result["number"]
+        number = result["countdown"]["number"]
         print(f"Number received: {number}")
 
         assert number == count
@@ -319,7 +334,7 @@ async def test_websocket_subscription_with_keepalive(
 
     async for result in session.subscribe(subscription):
 
-        number = result["number"]
+        number = result["countdown"]["number"]
         print(f"Number received: {number}")
 
         assert number == count
@@ -344,10 +359,54 @@ def test_websocket_subscription_sync(server, subscription_str):
 
     for result in client.subscribe(subscription):
 
-        number = result["number"]
+        number = result["countdown"]["number"]
         print(f"Number received: {number}")
 
         assert number == count
         count -= 1
 
     assert count == -1
+
+
+class NumberAddParser:
+    """ Class with a parse_value method used to increment a number """
+
+    def __init__(self, increment: int):
+        self.increment: int = increment
+
+    def parse_value(self, value: int) -> int:
+        return value + self.increment
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("server", [server_countdown], indirect=True)
+@pytest.mark.parametrize("subscription_str", [countdown_subscription_str])
+async def test_websocket_subscription_with_custom_types(
+    event_loop, server, subscription_str
+):
+
+    url = f"ws://{server.hostname}:{server.port}/graphql"
+
+    sample_transport = WebsocketsTransport(url=url)
+
+    count = 10
+    subscription = gql(subscription_str.format(count=count))
+
+    add_10 = NumberAddParser(10)
+
+    custom_types = {"Int": add_10}
+
+    # Instanciate a client which will add 10 to all the scalars of type Int received
+    async with Client(
+        transport=sample_transport, custom_types=custom_types, type_def=countdown_schema
+    ) as session:
+        async for result in session.subscribe(subscription):
+
+            number = result["countdown"]["number"]
+            print(f"Number received: {number}")
+
+            # We check here that the Int scalar has been correctly incremented by 10
+            assert number == count + 10
+            count -= 1
+
+        assert count == -1