Repository: mirumee/ariadne-codegen
Pull request #141 (2023-05-10 11:23:27)

feat: Add plugin to unwrap single field classes
Add a plugin that can be enabled by setting `unwrap_single_fields` to true in the settings to automatically return the inner value if a return type from a generated function only has a single method.

Ref: #66

---

This PR isn't in draft because the feature is complete but I would love some feedback on the code and approach. Even by following the suggestions in #66 there are probably things that could be solved different. I added a few TODOs/questions inline.

Also this PR doesn't yet include any tests but I saw you had a nice setup with `expected_client` for general client generated code so if this approach is resonable I'd be happy to ad some tests with the plugin enabled to ensure it's tested.

bombsimon: Thank you for the very quick and through review! I'll try to get back to addressing them as soon as possible.

Patch:
diff --git a/README.md b/README.md
index 6e4e53c..07cfa5a 100644
--- a/README.md
+++ b/README.md
@@ -85,6 +85,13 @@ Optional settings:
 Ariadne Codegen implements a plugin system that enables further customization and fine-tuning of generated Python code. It’s documentation is available separately in the [PLUGINS.md](https://github.com/mirumee/ariadne-codegen/blob/main/PLUGINS.md) file.
 
 
+### Standard plugins
+
+Ariadne Codegen ships with optional plugins importable from the `ariadne_codegen.contrib` package:
+
+- [`ariadne_codegen.contrib.shorter_results.ShorterResultsPlugin`](ariadne_codegen/contrib/shorter_results.py) - This plugin processes generated client methods for operations where only single top field is requested, so they return this field's value directly instead of operation's result type. For example get_user method generated for query `GetUser() { user(...) { ... }}` will return value of user field directly instead of `GetUserResult`.
+
+
 ## Using generated client
 
 Generated client can be imported from package:
@@ -133,7 +140,7 @@ If `type`/`serialize`/`parse` contains at least one `.` then string will be spli
 
 ### Example with scalar mapped to built-in type
 
-In this case scalar is mapped to built-in `str` which doesn\`t require custom `serialize ` and `parse` methods. 
+In this case scalar is mapped to built-in `str` which doesn't require custom `serialize ` and `parse` methods.
 
 ```toml
 [tool.ariadne-codegen.scalars.SCALARA]
@@ -189,7 +196,8 @@ This directive can be used along with `files_to_include` option to extend functi
 
 #### Example of usage of `mixin` and `files_to_include`:
 
-Query with `mixin` directive: 
+Query with `mixin` directive:
+
 ```gql
 query listUsers {
     users @mixin(from: ".mixins", import: "UsersMixin") {
@@ -245,7 +253,7 @@ Instead of generating client, you can generate file with a copy of GraphQL schem
 ariadne-codegen graphqlschema
 ```
 
-`graphqlschema` mode reads configuration from the same place as [`client`](#configuration) but uses only `schema_path`, `remote_schema_url`, `remote_schema_headers` and `remote_schema_verify_ssl` options with addition to some extra options specific to it:    
+`graphqlschema` mode reads configuration from the same place as [`client`](#configuration) but uses only `schema_path`, `remote_schema_url`, `remote_schema_headers` and `remote_schema_verify_ssl` options with addition to some extra options specific to it:
 
 - `target_file_path` (defaults to `"schema.py"`) - destination path for generated file
 - `schema_variable_name` (defaults to `"schema"`) - name for schema variable, must be valid python identifier
@@ -255,7 +263,7 @@ Generated file contains:
 
 - Necessary imports
 - Type map declaration `{type_map_variable_name}: TypeMap = {...}`
-- Schema declaration `{schema_variable_name}: GraphQLSchema = GraphQLSchema(...)` 
+- Schema declaration `{schema_variable_name}: GraphQLSchema = GraphQLSchema(...)`
 
 
 ## Contributing
diff --git a/ariadne_codegen/contrib/shorter_results.py b/ariadne_codegen/contrib/shorter_results.py
new file mode 100644
index 0000000..0c335eb
--- /dev/null
+++ b/ariadne_codegen/contrib/shorter_results.py
@@ -0,0 +1,356 @@
+"""
+Shorter results by returning single fields directly.
+
+Given a GraphQL schema that looks like this:
+
+    type Query{
+        me: User
+    }
+
+    type User {
+        id: ID!
+        username: String!
+    }
+
+The codegen will generate a class that has a single attribute `me`.
+
+    class GetAuthenticatedUser(BaseModel):
+        me: Optional["GetAuthenticatedUserMe"]
+
+When later using the client to call `get_authenticated_user` the return type
+would most of the time immediately be expanded to get the inner
+`GetAuthenticatedUserMe`:
+
+    me = (await client.get_authenticated_user).me
+
+By enabling this plugin the return type of the generated
+`get_authenticated_user` will instead directly return a
+`GetAuthenticatedUserMe` so the caller can use:
+
+    me = await client.get_authenticated_user
+
+This plugin can be enabled in the settings:
+
+    plugins = ["ariadne_codegen.contrib.shorter_results.ShorterResultsPlugin"]
+"""
+
+import ast
+from copy import deepcopy
+from typing import Dict, List, Optional, Union
+
+from graphql import ExecutableDefinitionNode, GraphQLSchema, SelectionSetNode
+
+from ..codegen import (
+    generate_async_for,
+    generate_attribute,
+    generate_expr,
+    generate_name,
+    generate_return,
+    generate_subscript,
+    generate_yield,
+)
+from ..plugins.base import Plugin
+
+
+class ShorterResultsPlugin(Plugin):
+    """
+    Make single field return types expanded.
+
+    All client methods that return a type with a single field will be expanded and
+    instead return the inner field.
+    """
+
+    def __init__(self, schema: GraphQLSchema, config_dict: Dict) -> None:
+        self.class_dict: Dict[str, ast.ClassDef] = {}
+        self.extended_imports: Dict[str, set] = {}
+
+        super().__init__(schema, config_dict)
+
+    def generate_result_class(
+        self,
+        class_def: ast.ClassDef,
+        operation_definition: ExecutableDefinitionNode,
+        selection_set: SelectionSetNode,
+    ) -> ast.ClassDef:
+        """Store a map of all classes and their AST."""
+        self.class_dict[class_def.name] = class_def
+
+        return super().generate_result_class(
+            class_def, operation_definition, selection_set
+        )
+
+    def generate_client_module(self, module: ast.Module) -> ast.Module:
+        """
+        Update the generated client.
+
+        At this point we've the full ast for all classes used for the client so
+        instead of modifying class by class in `generate_client_method` we
+        modify each method from here. By doing so we ensure that we have a map
+        between all classes, including fragments, so we can traverse them to
+        figure out how many fields we have in total for each request.
+
+        If we only have a single field, including when expanding our fragments,
+        we can expand the single field and update the return type.
+        """
+        client_def = next(
+            filter(lambda o: isinstance(o, ast.ClassDef), module.body), None
+        )
+        if not client_def or not isinstance(client_def, ast.ClassDef):
+            return super().generate_client_module(module)
+
+        for method_def in [
+            m
+            for m in client_def.body
+            if isinstance(m, (ast.FunctionDef, ast.AsyncFunctionDef))
+        ]:
+            self._modify_method_def(method_def)
+
+        if len(self.extended_imports) == 0:
+            return super().generate_client_module(module)
+
+        for stmt in module.body:
+            # We know we've already imported the wrapping type so we're only
+            # interested in `ImportFrom`.
+            if not isinstance(stmt, ast.ImportFrom):
+                continue
+
+            # The module we import from is always the same as the method we
+            # modified when changing the return type.
+            if stmt.module not in self.extended_imports:
+                continue
+
+            # Add all additional imports discovered when generating the client
+            # method.
+            for additional_import in self.extended_imports[stmt.module]:
+                stmt.names.append(ast.alias(name=additional_import))
+
+        return super().generate_client_module(module)
+
+    def _modify_method_def(
+        self, method_def: Union[ast.FunctionDef, ast.AsyncFunctionDef]
+    ) -> None:
+        """
+        Change the method generated in the client to call the inner field and
+        change the return type. We do this here instead of
+        `generate_client_method` to ensure we have a map between all our classes
+        when updating the method def.
+        """
+        if len(method_def.body) < 1:
+            return None
+
+        last_stmt = method_def.body[-1]
+
+        if isinstance(last_stmt, ast.Return):
+            self._generate_query_and_mutation_client_method(method_def, last_stmt)
+        elif isinstance(last_stmt, ast.AsyncFor):
+            self._generate_subscription_client_method(method_def, last_stmt)
+
+        return None
+
+    def _generate_subscription_client_method(
+        self,
+        method_def: Union[ast.FunctionDef, ast.AsyncFunctionDef],
+        async_for_stmt: ast.AsyncFor,
+    ):
+        """Generate method for subscriptions."""
+        if not isinstance(method_def.returns, ast.Subscript):
+            return None
+
+        if not isinstance(method_def.returns.slice, ast.Name):
+            return None
+
+        node_and_class = _return_or_yield_node_and_class(
+            method_def.returns.slice.id, self.class_dict
+        )
+        if node_and_class is None:
+            return None
+
+        yield_node, single_field_classes, single_field_return_class = node_and_class
+
+        previous_yield_value = _get_yield_value_from_async_for(method_def.body[-1])
+        if previous_yield_value is None:
+            return None
+
+        # Update the return type to be the inner field and ensure we reference
+        # the single field before returning.
+        method_def.returns = generate_subscript(
+            value=generate_name("AsyncIterator"),
+            slice_=yield_node,
+        )
+
+        method_def.body[-1] = generate_async_for(
+            target=async_for_stmt.target,
+            iter_=async_for_stmt.iter,
+            body=generate_expr(
+                value=generate_yield(
+                    value=generate_attribute(
+                        value=previous_yield_value,
+                        attr=single_field_return_class,
+                    )
+                )
+            ),
+        )
+
+        self._update_imports(method_def, single_field_classes)
+
+        return None
+
+    def _generate_query_and_mutation_client_method(
+        self,
+        method_def: Union[ast.FunctionDef, ast.AsyncFunctionDef],
+        return_stmt: ast.Return,
+    ) -> None:
+        """Generate method for query or mutations."""
+        if return_stmt.value is None:
+            return None
+
+        if not isinstance(method_def.returns, ast.Name):
+            return None
+
+        node_and_class = _return_or_yield_node_and_class(
+            method_def.returns.id, self.class_dict
+        )
+        if node_and_class is None:
+            return None
+
+        return_node, single_field_classes, single_field_return_class = node_and_class
+
+        # Update the return type to be the inner field and ensure we reference
+        # the single field before returning.
+        method_def.returns = return_node
+        method_def.body[-1] = generate_return(
+            value=generate_attribute(
+                value=return_stmt.value, attr=single_field_return_class
+            )
+        )
+
+        self._update_imports(method_def, single_field_classes)
+
+        return None
+
+    def _update_imports(
+        self,
+        method_def: Union[ast.FunctionDef, ast.AsyncFunctionDef],
+        single_field_classes: List[str],
+    ):
+        """
+        After expanding to the inner type we will end up with one or more
+        classes that we're now returning. Ensure they're all imported in the
+        file defining the method.
+        """
+        for single_field_class in single_field_classes:
+            if single_field_class not in self.class_dict:
+                continue
+
+            # After we change the type we also need to import it in the client if
+            # it's one of our generated types so add the extra import as needed.
+            if method_def.name not in self.extended_imports:
+                self.extended_imports[method_def.name] = set()
+
+            self.extended_imports[method_def.name].add(single_field_class)
+
+
+def _get_yield_value_from_async_for(stmt: ast.stmt) -> Optional[ast.expr]:
+    """
+    Extract inner yield value from `ast.AsyncFor` and return `None` if the stmt
+    isn't of expected type.
+    """
+    if not isinstance(stmt, ast.AsyncFor):
+        return None
+
+    if len(stmt.body) < 1:
+        return None
+
+    if not isinstance(stmt.body[0], ast.Expr):
+        return None
+
+    if not isinstance(stmt.body[0].value, ast.Yield):
+        return None
+
+    return stmt.body[0].value.value
+
+
+def _return_or_yield_node_and_class(
+    current_return_class: str, class_dict: Dict[str, ast.ClassDef]
+) -> Optional[tuple[ast.expr, List[str], str]]:
+    """
+    Given there's only a single field in the return type, return the ast for
+    that node, what class(es) that is and the name of the field.
+    """
+    if current_return_class not in class_dict:
+        return None
+
+    fields = _get_all_fields(class_dict[current_return_class], class_dict)
+    if len(fields) != 1:
+        return None
+
+    single_field = fields[0]
+    if not isinstance(single_field, ast.AnnAssign):
+        return None
+
+    if not isinstance(single_field.target, ast.Name):
+        return None
+
+    # Traverse the type annotation until we find the inner type. This can
+    # require several iterations if the return type is something like
+    # Optional[List[Any]]
+    #
+    # We make a deepcopy because we need to keep the quoted annotations for the
+    # fields in the query classes but we want to have it unquoted in the client
+    # method where we import the actual types.
+    annotations = deepcopy(single_field.annotation)
+    return_node, return_classes = _update_node(annotations)
+
+    return return_node, return_classes, single_field.target.id
+
+
+def _update_node(node: ast.expr) -> tuple[ast.expr, List[str]]:
+    """
+    Walk down a node to find the inner `ast.Name`. Once found, evaluate the
+    potential literal so it gets unquoted and return the inner name.
+    """
+    if isinstance(node, ast.Name):
+        try:
+            node.id = ast.literal_eval(node.id)
+        except ValueError:
+            # Not a literal
+            pass
+
+        return node, [node.id]
+
+    if isinstance(node, ast.Subscript):
+        node.slice, node_ids = _update_node(node.slice)
+
+        return node, node_ids
+
+    if isinstance(node, ast.Tuple):
+        all_node_ids = []
+        for i, elt in enumerate(node.elts):
+            node.elts[i], node_ids = _update_node(elt)
+            all_node_ids.extend(node_ids)
+
+        return node, all_node_ids
+
+    return ast.expr(), []
+
+
+def _get_all_fields(
+    class_def: ast.ClassDef, class_dict: Dict[str, ast.ClassDef]
+) -> List[ast.AnnAssign]:
+    """
+    Recursively get all fields from all inherited classes to figure out the
+    total number of fields.
+    """
+    fields = []
+
+    for base in class_def.bases:
+        if not isinstance(base, ast.Name) or base.id not in class_dict:
+            continue
+
+        fields.extend(_get_all_fields(class_dict[base.id], class_dict))
+
+    for field in class_def.body:
+        if isinstance(field, ast.AnnAssign):
+            fields.append(field)
+
+    return fields


Test patch:
diff --git a/tests/main/clients/shorter_results/expected_client/__init__.py b/tests/main/clients/shorter_results/expected_client/__init__.py
new file mode 100644
index 0000000..8c00408
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/__init__.py
@@ -0,0 +1,53 @@
+from .async_base_client import AsyncBaseClient
+from .base_model import BaseModel
+from .client import Client
+from .exceptions import (
+    GraphQLClientError,
+    GraphQLClientGraphQLError,
+    GraphQLClientGraphQLMultiError,
+    GraphQLClientHttpError,
+    GraphQlClientInvalidResponseError,
+)
+from .fragments import AnimalFragment
+from .get_animal_by_name import (
+    GetAnimalByName,
+    GetAnimalByNameAnimalByNameAnimal,
+    GetAnimalByNameAnimalByNameCat,
+    GetAnimalByNameAnimalByNameDog,
+)
+from .get_animal_fragment import GetAnimalFragment
+from .get_animal_fragment_with_extra import GetAnimalFragmentWithExtra
+from .get_authenticated_user import GetAuthenticatedUser, GetAuthenticatedUserMe
+from .list_strings_1 import ListStrings1
+from .list_strings_2 import ListStrings2
+from .list_strings_3 import ListStrings3
+from .list_strings_4 import ListStrings4
+from .list_type_a import ListTypeA, ListTypeAListOptionalTypeA
+from .subscribe_strings import SubscribeStrings
+
+__all__ = [
+    "AnimalFragment",
+    "AsyncBaseClient",
+    "BaseModel",
+    "Client",
+    "GetAnimalByName",
+    "GetAnimalByNameAnimalByNameAnimal",
+    "GetAnimalByNameAnimalByNameCat",
+    "GetAnimalByNameAnimalByNameDog",
+    "GetAnimalFragment",
+    "GetAnimalFragmentWithExtra",
+    "GetAuthenticatedUser",
+    "GetAuthenticatedUserMe",
+    "GraphQLClientError",
+    "GraphQLClientGraphQLError",
+    "GraphQLClientGraphQLMultiError",
+    "GraphQLClientHttpError",
+    "GraphQlClientInvalidResponseError",
+    "ListStrings1",
+    "ListStrings2",
+    "ListStrings3",
+    "ListStrings4",
+    "ListTypeA",
+    "ListTypeAListOptionalTypeA",
+    "SubscribeStrings",
+]
diff --git a/tests/main/clients/shorter_results/expected_client/async_base_client.py b/tests/main/clients/shorter_results/expected_client/async_base_client.py
new file mode 100644
index 0000000..86f3e06
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/async_base_client.py
@@ -0,0 +1,210 @@
+import enum
+import json
+from typing import Any, AsyncIterator, Dict, Optional, TypeVar, cast
+from uuid import uuid4
+
+import httpx
+from pydantic import BaseModel
+
+from .base_model import UNSET
+from .exceptions import (
+    GraphQLClientGraphQLMultiError,
+    GraphQLClientHttpError,
+    GraphQLClientInvalidMessageFormat,
+    GraphQlClientInvalidResponseError,
+)
+
+try:
+    from websockets.client import WebSocketClientProtocol
+    from websockets.client import connect as ws_connect
+    from websockets.typing import Data, Origin, Subprotocol
+except ImportError:
+    from contextlib import asynccontextmanager
+
+    @asynccontextmanager  # type: ignore
+    async def ws_connect(*args, **kwargs):  # pylint: disable=unused-argument
+        raise NotImplementedError("Subscriptions require 'websockets' package.")
+        yield  # pylint: disable=unreachable
+
+    WebSocketClientProtocol = Any  # type: ignore
+    Data = Any  # type: ignore
+    Origin = Any  # type: ignore
+    Subprotocol = Any  # type: ignore
+
+
+Self = TypeVar("Self", bound="AsyncBaseClient")
+
+GRAPHQL_TRANSPORT_WS = "graphql-transport-ws"
+
+
+class GraphQLTransportWSMessageType(str, enum.Enum):
+    CONNECTION_INIT = "connection_init"
+    CONNECTION_ACK = "connection_ack"
+    PING = "ping"
+    PONG = "pong"
+    SUBSCRIBE = "subscribe"
+    NEXT = "next"
+    ERROR = "error"
+    COMPLETE = "complete"
+
+
+class AsyncBaseClient:
+    def __init__(
+        self,
+        url: str = "",
+        headers: Optional[Dict[str, str]] = None,
+        http_client: Optional[httpx.AsyncClient] = None,
+        ws_url: str = "",
+        ws_headers: Optional[Dict[str, Any]] = None,
+        ws_origin: Optional[str] = None,
+        ws_connection_init_payload: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self.url = url
+        self.headers = headers
+        self.http_client = (
+            http_client if http_client else httpx.AsyncClient(headers=headers)
+        )
+        self.ws_url = ws_url
+        self.ws_headers = ws_headers or {}
+        self.ws_origin = Origin(ws_origin) if ws_origin else None
+        self.ws_connection_init_payload = ws_connection_init_payload
+
+    async def __aenter__(self: Self) -> Self:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: object,
+        exc_val: object,
+        exc_tb: object,
+    ) -> None:
+        await self.http_client.aclose()
+
+    async def execute(
+        self, query: str, variables: Optional[Dict[str, Any]] = None
+    ) -> httpx.Response:
+        payload: Dict[str, Any] = {"query": query}
+        if variables:
+            payload["variables"] = self._convert_dict_to_json_serializable(variables)
+        return await self.http_client.post(url=self.url, json=payload)
+
+    def get_data(self, response: httpx.Response) -> Dict[str, Any]:
+        if not response.is_success:
+            raise GraphQLClientHttpError(
+                status_code=response.status_code, response=response
+            )
+
+        try:
+            response_json = response.json()
+        except ValueError as exc:
+            raise GraphQlClientInvalidResponseError(response=response) from exc
+
+        if (not isinstance(response_json, dict)) or ("data" not in response_json):
+            raise GraphQlClientInvalidResponseError(response=response)
+
+        data = response_json["data"]
+        errors = response_json.get("errors")
+
+        if errors:
+            raise GraphQLClientGraphQLMultiError.from_errors_dicts(
+                errors_dicts=errors, data=data
+            )
+
+        return cast(dict[str, Any], data)
+
+    async def execute_ws(
+        self, query: str, variables: Optional[Dict[str, Any]] = None
+    ) -> AsyncIterator[Dict[str, Any]]:
+        operation_id = str(uuid4())
+        async with ws_connect(
+            self.ws_url,
+            subprotocols=[Subprotocol(GRAPHQL_TRANSPORT_WS)],
+            origin=self.ws_origin,
+            extra_headers=self.ws_headers,
+        ) as websocket:
+            await self._send_connection_init(websocket)
+            await self._send_subscribe(
+                websocket,
+                operation_id=operation_id,
+                query=query,
+                variables=variables,
+            )
+
+            async for message in websocket:
+                data = await self._handle_ws_message(message, websocket)
+                if data:
+                    yield data
+
+    def _convert_dict_to_json_serializable(
+        self, dict_: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        return {
+            key: self._convert_value(value)
+            for key, value in dict_.items()
+            if value is not UNSET
+        }
+
+    def _convert_value(self, value: Any) -> Any:
+        if isinstance(value, BaseModel):
+            return value.dict(by_alias=True, exclude_unset=True)
+        if isinstance(value, list):
+            return [self._convert_value(item) for item in value]
+        return value
+
+    async def _send_connection_init(self, websocket: WebSocketClientProtocol) -> None:
+        payload: Dict[str, Any] = {
+            "type": GraphQLTransportWSMessageType.CONNECTION_INIT.value
+        }
+        if self.ws_connection_init_payload:
+            payload["payload"] = self.ws_connection_init_payload
+        await websocket.send(json.dumps(payload))
+
+    async def _send_subscribe(
+        self,
+        websocket: WebSocketClientProtocol,
+        operation_id: str,
+        query: str,
+        variables: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        payload: Dict[str, Any] = {
+            "id": operation_id,
+            "type": GraphQLTransportWSMessageType.SUBSCRIBE.value,
+            "payload": {"query": query},
+        }
+        if variables:
+            payload["payload"]["variables"] = self._convert_dict_to_json_serializable(
+                variables
+            )
+        await websocket.send(json.dumps(payload))
+
+    async def _handle_ws_message(
+        self, message: Data, websocket: WebSocketClientProtocol
+    ) -> Optional[Dict[str, Any]]:
+        try:
+            message_dict = json.loads(message)
+        except json.JSONDecodeError as exc:
+            raise GraphQLClientInvalidMessageFormat(message=message) from exc
+
+        type_ = message_dict.get("type")
+        payload = message_dict.get("payload", {})
+
+        if not type_ or type_ not in {t.value for t in GraphQLTransportWSMessageType}:
+            raise GraphQLClientInvalidMessageFormat(message=message)
+
+        if type_ == GraphQLTransportWSMessageType.NEXT:
+            if "data" not in payload:
+                raise GraphQLClientInvalidMessageFormat(message=message)
+            return cast(Dict[str, Any], payload["data"])
+
+        if type_ == GraphQLTransportWSMessageType.COMPLETE:
+            await websocket.close()
+        elif type_ == GraphQLTransportWSMessageType.PING:
+            await websocket.send(
+                json.dumps({"type": GraphQLTransportWSMessageType.PONG.value})
+            )
+        elif type_ == GraphQLTransportWSMessageType.ERROR:
+            raise GraphQLClientGraphQLMultiError.from_errors_dicts(
+                errors_dicts=payload, data=message_dict
+            )
+
+        return None
diff --git a/tests/main/clients/shorter_results/expected_client/base_model.py b/tests/main/clients/shorter_results/expected_client/base_model.py
new file mode 100644
index 0000000..b1ec15e
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/base_model.py
@@ -0,0 +1,58 @@
+from typing import Any, Dict, Type, Union, get_args, get_origin
+
+from pydantic import BaseModel as PydanticBaseModel
+from pydantic.class_validators import validator
+from pydantic.fields import ModelField
+
+from .scalars import SCALARS_PARSE_FUNCTIONS, SCALARS_SERIALIZE_FUNCTIONS
+
+
+class UnsetType:
+    def __bool__(self) -> bool:
+        return False
+
+
+UNSET = UnsetType()
+
+
+class BaseModel(PydanticBaseModel):
+    class Config:
+        allow_population_by_field_name = True
+        validate_assignment = True
+        arbitrary_types_allowed = True
+
+    # pylint: disable=no-self-argument
+    @validator("*", pre=True)
+    def parse_custom_scalars(cls, value: Any, field: ModelField) -> Any:
+        return cls._parse_custom_scalar_value(value, field.annotation)
+
+    @classmethod
+    def _parse_custom_scalar_value(cls, value: Any, type_: Type[Any]) -> Any:
+        origin = get_origin(type_)
+        args = get_args(type_)
+        if origin is list and isinstance(value, list):
+            return [cls._parse_custom_scalar_value(item, args[0]) for item in value]
+
+        if origin is Union and type(None) in args:
+            sub_type: Any = list(filter(None, args))[0]
+            return cls._parse_custom_scalar_value(value, sub_type)
+
+        decode = SCALARS_PARSE_FUNCTIONS.get(type_)
+        if value and decode and callable(decode):
+            return decode(value)
+
+        return value
+
+    def dict(self, **kwargs: Any) -> Dict[str, Any]:
+        dict_ = super().dict(**kwargs)
+        return {key: self._serialize_value(value) for key, value in dict_.items()}
+
+    def _serialize_value(self, value: Any) -> Any:
+        serialize = SCALARS_SERIALIZE_FUNCTIONS.get(type(value))
+        if serialize and callable(serialize):
+            return serialize(value)
+
+        if isinstance(value, list):
+            return [self._serialize_value(item) for item in value]
+
+        return value
diff --git a/tests/main/clients/shorter_results/expected_client/client.py b/tests/main/clients/shorter_results/expected_client/client.py
new file mode 100644
index 0000000..6e73b20
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/client.py
@@ -0,0 +1,182 @@
+from typing import AsyncIterator, List, Optional, Union
+
+from .async_base_client import AsyncBaseClient
+from .get_animal_by_name import (
+    GetAnimalByName,
+    GetAnimalByNameAnimalByNameAnimal,
+    GetAnimalByNameAnimalByNameCat,
+    GetAnimalByNameAnimalByNameDog,
+)
+from .get_animal_fragment import GetAnimalFragment
+from .get_animal_fragment_with_extra import GetAnimalFragmentWithExtra
+from .get_authenticated_user import GetAuthenticatedUser, GetAuthenticatedUserMe
+from .list_strings_1 import ListStrings1
+from .list_strings_2 import ListStrings2
+from .list_strings_3 import ListStrings3
+from .list_strings_4 import ListStrings4
+from .list_type_a import ListTypeA, ListTypeAListOptionalTypeA
+from .subscribe_strings import SubscribeStrings
+
+
+def gql(q: str) -> str:
+    return q
+
+
+class Client(AsyncBaseClient):
+    async def get_authenticated_user(self) -> GetAuthenticatedUserMe:
+        query = gql(
+            """
+            query GetAuthenticatedUser {
+              me {
+                id
+                username
+              }
+            }
+            """
+        )
+        variables: dict[str, object] = {}
+        response = await self.execute(query=query, variables=variables)
+        data = self.get_data(response)
+        return GetAuthenticatedUser.parse_obj(data).me
+
+    async def list_strings_1(self) -> Optional[List[Optional[str]]]:
+        query = gql(
+            """
+            query ListStrings_1 {
+              optionalListOptionalString
+            }
+            """
+        )
+        variables: dict[str, object] = {}
+        response = await self.execute(query=query, variables=variables)
+        data = self.get_data(response)
+        return ListStrings1.parse_obj(data).optional_list_optional_string
+
+    async def list_strings_2(self) -> Optional[List[str]]:
+        query = gql(
+            """
+            query ListStrings_2 {
+              optionalListString
+            }
+            """
+        )
+        variables: dict[str, object] = {}
+        response = await self.execute(query=query, variables=variables)
+        data = self.get_data(response)
+        return ListStrings2.parse_obj(data).optional_list_string
+
+    async def list_strings_3(self) -> List[Optional[str]]:
+        query = gql(
+            """
+            query ListStrings_3 {
+              listOptionalString
+            }
+            """
+        )
+        variables: dict[str, object] = {}
+        response = await self.execute(query=query, variables=variables)
+        data = self.get_data(response)
+        return ListStrings3.parse_obj(data).list_optional_string
+
+    async def list_strings_4(self) -> List[str]:
+        query = gql(
+            """
+            query ListStrings_4 {
+              listString
+            }
+            """
+        )
+        variables: dict[str, object] = {}
+        response = await self.execute(query=query, variables=variables)
+        data = self.get_data(response)
+        return ListStrings4.parse_obj(data).list_string
+
+    async def list_type_a(self) -> List[Optional[ListTypeAListOptionalTypeA]]:
+        query = gql(
+            """
+            query ListTypeA {
+              listOptionalTypeA {
+                id
+              }
+            }
+            """
+        )
+        variables: dict[str, object] = {}
+        response = await self.execute(query=query, variables=variables)
+        data = self.get_data(response)
+        return ListTypeA.parse_obj(data).list_optional_type_a
+
+    async def get_animal_by_name(
+        self,
+    ) -> Union[
+        GetAnimalByNameAnimalByNameAnimal,
+        GetAnimalByNameAnimalByNameCat,
+        GetAnimalByNameAnimalByNameDog,
+    ]:
+        query = gql(
+            """
+            query GetAnimalByName {
+              animalByName {
+                __typename
+                name
+                ... on Cat {
+                  kittens
+                }
+                ... on Dog {
+                  puppies
+                }
+              }
+            }
+            """
+        )
+        variables: dict[str, object] = {}
+        response = await self.execute(query=query, variables=variables)
+        data = self.get_data(response)
+        return GetAnimalByName.parse_obj(data).animal_by_name
+
+    async def get_animal_fragment(self) -> str:
+        query = gql(
+            """
+            query GetAnimalFragment {
+              ...AnimalFragment
+            }
+
+            fragment AnimalFragment on Animal {
+              name
+            }
+            """
+        )
+        variables: dict[str, object] = {}
+        response = await self.execute(query=query, variables=variables)
+        data = self.get_data(response)
+        return GetAnimalFragment.parse_obj(data).name
+
+    async def get_animal_fragment_with_extra(self) -> GetAnimalFragmentWithExtra:
+        query = gql(
+            """
+            query GetAnimalFragmentWithExtra {
+              ...AnimalFragment
+              listString
+            }
+
+            fragment AnimalFragment on Animal {
+              name
+            }
+            """
+        )
+        variables: dict[str, object] = {}
+        response = await self.execute(query=query, variables=variables)
+        data = self.get_data(response)
+        return GetAnimalFragmentWithExtra.parse_obj(data)
+
+    async def subscribe_strings(self) -> AsyncIterator[Optional[List[str]]]:
+        query = gql(
+            """
+            subscription SubscribeStrings {
+              optionalListString
+            }
+            """
+        )
+        variables: dict[str, object] = {}
+        async for data in self.execute_ws(query=query, variables=variables):
+            yield SubscribeStrings.parse_obj(data).optional_list_string
diff --git a/tests/main/clients/shorter_results/expected_client/enums.py b/tests/main/clients/shorter_results/expected_client/enums.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/main/clients/shorter_results/expected_client/exceptions.py b/tests/main/clients/shorter_results/expected_client/exceptions.py
new file mode 100644
index 0000000..830a44d
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/exceptions.py
@@ -0,0 +1,79 @@
+from typing import Any, Dict, List, Optional, Union
+
+import httpx
+
+
+class GraphQLClientError(Exception):
+    """Base exception."""
+
+
+class GraphQLClientHttpError(GraphQLClientError):
+    def __init__(self, status_code: int, response: httpx.Response) -> None:
+        self.status_code = status_code
+        self.response = response
+
+    def __str__(self) -> str:
+        return f"HTTP status code: {self.status_code}"
+
+
+class GraphQlClientInvalidResponseError(GraphQLClientError):
+    def __init__(self, response: httpx.Response) -> None:
+        self.response = response
+
+    def __str__(self) -> str:
+        return "Invalid response format."
+
+
+class GraphQLClientGraphQLError(GraphQLClientError):
+    def __init__(
+        self,
+        message: str,
+        locations: Optional[List[Dict[str, int]]] = None,
+        path: Optional[List[str]] = None,
+        extensions: Optional[Dict[str, object]] = None,
+        orginal: Optional[Dict[str, object]] = None,
+    ):
+        self.message = message
+        self.locations = locations
+        self.path = path
+        self.extensions = extensions
+        self.orginal = orginal
+
+    def __str__(self) -> str:
+        return self.message
+
+    @classmethod
+    def from_dict(cls, error: dict[str, Any]) -> "GraphQLClientGraphQLError":
+        return cls(
+            message=error["message"],
+            locations=error.get("locations"),
+            path=error.get("path"),
+            extensions=error.get("extensions"),
+            orginal=error,
+        )
+
+
+class GraphQLClientGraphQLMultiError(GraphQLClientError):
+    def __init__(self, errors: List[GraphQLClientGraphQLError], data: dict[str, Any]):
+        self.errors = errors
+        self.data = data
+
+    def __str__(self) -> str:
+        return "; ".join(str(e) for e in self.errors)
+
+    @classmethod
+    def from_errors_dicts(
+        cls, errors_dicts: List[dict[str, Any]], data: dict[str, Any]
+    ) -> "GraphQLClientGraphQLMultiError":
+        return cls(
+            errors=[GraphQLClientGraphQLError.from_dict(e) for e in errors_dicts],
+            data=data,
+        )
+
+
+class GraphQLClientInvalidMessageFormat(GraphQLClientError):
+    def __init__(self, message: Union[str, bytes]) -> None:
+        self.message = message
+
+    def __str__(self) -> str:
+        return "Invalid message format."
diff --git a/tests/main/clients/shorter_results/expected_client/fragments.py b/tests/main/clients/shorter_results/expected_client/fragments.py
new file mode 100644
index 0000000..c7ac556
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/fragments.py
@@ -0,0 +1,8 @@
+from .base_model import BaseModel
+
+
+class AnimalFragment(BaseModel):
+    name: str
+
+
+AnimalFragment.update_forward_refs()
diff --git a/tests/main/clients/shorter_results/expected_client/get_animal_by_name.py b/tests/main/clients/shorter_results/expected_client/get_animal_by_name.py
new file mode 100644
index 0000000..4e43250
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/get_animal_by_name.py
@@ -0,0 +1,36 @@
+from typing import Union
+
+from pydantic import Field
+
+from .base_model import BaseModel
+
+
+class GetAnimalByName(BaseModel):
+    animal_by_name: Union[
+        "GetAnimalByNameAnimalByNameAnimal",
+        "GetAnimalByNameAnimalByNameCat",
+        "GetAnimalByNameAnimalByNameDog",
+    ] = Field(alias="animalByName")
+
+
+class GetAnimalByNameAnimalByNameAnimal(BaseModel):
+    __typename__: str = Field(alias="__typename")
+    name: str
+
+
+class GetAnimalByNameAnimalByNameCat(BaseModel):
+    __typename__: str = Field(alias="__typename")
+    name: str
+    kittens: int
+
+
+class GetAnimalByNameAnimalByNameDog(BaseModel):
+    __typename__: str = Field(alias="__typename")
+    name: str
+    puppies: int
+
+
+GetAnimalByName.update_forward_refs()
+GetAnimalByNameAnimalByNameAnimal.update_forward_refs()
+GetAnimalByNameAnimalByNameCat.update_forward_refs()
+GetAnimalByNameAnimalByNameDog.update_forward_refs()
diff --git a/tests/main/clients/shorter_results/expected_client/get_animal_fragment.py b/tests/main/clients/shorter_results/expected_client/get_animal_fragment.py
new file mode 100644
index 0000000..fb6bb9b
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/get_animal_fragment.py
@@ -0,0 +1,8 @@
+from .fragments import AnimalFragment
+
+
+class GetAnimalFragment(AnimalFragment):
+    pass
+
+
+GetAnimalFragment.update_forward_refs()
diff --git a/tests/main/clients/shorter_results/expected_client/get_animal_fragment_with_extra.py b/tests/main/clients/shorter_results/expected_client/get_animal_fragment_with_extra.py
new file mode 100644
index 0000000..9133a97
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/get_animal_fragment_with_extra.py
@@ -0,0 +1,12 @@
+from typing import List
+
+from pydantic import Field
+
+from .fragments import AnimalFragment
+
+
+class GetAnimalFragmentWithExtra(AnimalFragment):
+    list_string: List[str] = Field(alias="listString")
+
+
+GetAnimalFragmentWithExtra.update_forward_refs()
diff --git a/tests/main/clients/shorter_results/expected_client/get_authenticated_user.py b/tests/main/clients/shorter_results/expected_client/get_authenticated_user.py
new file mode 100644
index 0000000..e856030
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/get_authenticated_user.py
@@ -0,0 +1,14 @@
+from .base_model import BaseModel
+
+
+class GetAuthenticatedUser(BaseModel):
+    me: "GetAuthenticatedUserMe"
+
+
+class GetAuthenticatedUserMe(BaseModel):
+    id: str
+    username: str
+
+
+GetAuthenticatedUser.update_forward_refs()
+GetAuthenticatedUserMe.update_forward_refs()
diff --git a/tests/main/clients/shorter_results/expected_client/input_types.py b/tests/main/clients/shorter_results/expected_client/input_types.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/main/clients/shorter_results/expected_client/list_strings_1.py b/tests/main/clients/shorter_results/expected_client/list_strings_1.py
new file mode 100644
index 0000000..7f87dad
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/list_strings_1.py
@@ -0,0 +1,14 @@
+from typing import List, Optional
+
+from pydantic import Field
+
+from .base_model import BaseModel
+
+
+class ListStrings1(BaseModel):
+    optional_list_optional_string: Optional[List[Optional[str]]] = Field(
+        alias="optionalListOptionalString"
+    )
+
+
+ListStrings1.update_forward_refs()
diff --git a/tests/main/clients/shorter_results/expected_client/list_strings_2.py b/tests/main/clients/shorter_results/expected_client/list_strings_2.py
new file mode 100644
index 0000000..0467eec
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/list_strings_2.py
@@ -0,0 +1,12 @@
+from typing import List, Optional
+
+from pydantic import Field
+
+from .base_model import BaseModel
+
+
+class ListStrings2(BaseModel):
+    optional_list_string: Optional[List[str]] = Field(alias="optionalListString")
+
+
+ListStrings2.update_forward_refs()
diff --git a/tests/main/clients/shorter_results/expected_client/list_strings_3.py b/tests/main/clients/shorter_results/expected_client/list_strings_3.py
new file mode 100644
index 0000000..567ab67
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/list_strings_3.py
@@ -0,0 +1,12 @@
+from typing import List, Optional
+
+from pydantic import Field
+
+from .base_model import BaseModel
+
+
+class ListStrings3(BaseModel):
+    list_optional_string: List[Optional[str]] = Field(alias="listOptionalString")
+
+
+ListStrings3.update_forward_refs()
diff --git a/tests/main/clients/shorter_results/expected_client/list_strings_4.py b/tests/main/clients/shorter_results/expected_client/list_strings_4.py
new file mode 100644
index 0000000..304a116
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/list_strings_4.py
@@ -0,0 +1,12 @@
+from typing import List
+
+from pydantic import Field
+
+from .base_model import BaseModel
+
+
+class ListStrings4(BaseModel):
+    list_string: List[str] = Field(alias="listString")
+
+
+ListStrings4.update_forward_refs()
diff --git a/tests/main/clients/shorter_results/expected_client/list_type_a.py b/tests/main/clients/shorter_results/expected_client/list_type_a.py
new file mode 100644
index 0000000..5baca05
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/list_type_a.py
@@ -0,0 +1,19 @@
+from typing import List, Optional
+
+from pydantic import Field
+
+from .base_model import BaseModel
+
+
+class ListTypeA(BaseModel):
+    list_optional_type_a: List[Optional["ListTypeAListOptionalTypeA"]] = Field(
+        alias="listOptionalTypeA"
+    )
+
+
+class ListTypeAListOptionalTypeA(BaseModel):
+    id: int
+
+
+ListTypeA.update_forward_refs()
+ListTypeAListOptionalTypeA.update_forward_refs()
diff --git a/tests/main/clients/shorter_results/expected_client/scalars.py b/tests/main/clients/shorter_results/expected_client/scalars.py
new file mode 100644
index 0000000..568ad60
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/scalars.py
@@ -0,0 +1,4 @@
+from typing import Any, Callable, Dict
+
+SCALARS_PARSE_FUNCTIONS: Dict[Any, Callable[[str], Any]] = {}
+SCALARS_SERIALIZE_FUNCTIONS: Dict[Any, Callable[[Any], str]] = {}
diff --git a/tests/main/clients/shorter_results/expected_client/subscribe_strings.py b/tests/main/clients/shorter_results/expected_client/subscribe_strings.py
new file mode 100644
index 0000000..f1b5e02
--- /dev/null
+++ b/tests/main/clients/shorter_results/expected_client/subscribe_strings.py
@@ -0,0 +1,12 @@
+from typing import List, Optional
+
+from pydantic import Field
+
+from .base_model import BaseModel
+
+
+class SubscribeStrings(BaseModel):
+    optional_list_string: Optional[List[str]] = Field(alias="optionalListString")
+
+
+SubscribeStrings.update_forward_refs()
diff --git a/tests/main/clients/shorter_results/pyproject.toml b/tests/main/clients/shorter_results/pyproject.toml
new file mode 100644
index 0000000..24d3eaf
--- /dev/null
+++ b/tests/main/clients/shorter_results/pyproject.toml
@@ -0,0 +1,6 @@
+[tool.ariadne-codegen]
+schema_path = "schema.graphql"
+queries_path = "queries.graphql"
+include_comments = false
+target_package_name = "shorter_results"
+plugins = ["ariadne_codegen.contrib.shorter_results.ShorterResultsPlugin"]
diff --git a/tests/main/clients/shorter_results/queries.graphql b/tests/main/clients/shorter_results/queries.graphql
new file mode 100644
index 0000000..c4f3cc4
--- /dev/null
+++ b/tests/main/clients/shorter_results/queries.graphql
@@ -0,0 +1,57 @@
+query GetAuthenticatedUser {
+  me {
+    id
+    username
+  }
+}
+
+query ListStrings_1 {
+  optionalListOptionalString
+}
+
+query ListStrings_2 {
+  optionalListString
+}
+
+query ListStrings_3 {
+  listOptionalString
+}
+
+query ListStrings_4 {
+  listString
+}
+
+query ListTypeA {
+  listOptionalTypeA {
+    id
+  }
+}
+
+query GetAnimalByName {
+  animalByName {
+    name
+    ... on Cat {
+      kittens
+    }
+    ... on Dog {
+      puppies
+    }
+  }
+}
+
+query GetAnimalFragment {
+  ...AnimalFragment
+}
+
+query GetAnimalFragmentWithExtra {
+  ...AnimalFragment
+  listString
+}
+
+fragment AnimalFragment on Animal {
+  name
+}
+
+subscription SubscribeStrings {
+  optionalListString
+}
diff --git a/tests/main/clients/shorter_results/schema.graphql b/tests/main/clients/shorter_results/schema.graphql
new file mode 100644
index 0000000..dcc9437
--- /dev/null
+++ b/tests/main/clients/shorter_results/schema.graphql
@@ -0,0 +1,39 @@
+type TypeA {
+  id: Int!
+}
+
+type User {
+    id: ID!
+    username: String!
+}
+
+interface Animal {
+  name: String!
+}
+
+type Cat implements Animal {
+  name: String!
+  kittens: Int!
+}
+
+type Dog implements Animal {
+  name: String!
+  puppies: Int!
+}
+
+type Query {
+  optionalListOptionalString: [String]
+  optionalListString: [String!]
+  listOptionalString: [String]!
+  listString: [String!]!
+
+  listOptionalTypeA: [TypeA]!
+
+  me: User!
+
+  animalByName(name: String!): Animal!
+}
+
+type Subscription {
+  optionalListString: [String!]
+}
diff --git a/tests/main/test_main.py b/tests/main/test_main.py
index 59e85c9..8a9129a 100644
--- a/tests/main/test_main.py
+++ b/tests/main/test_main.py
@@ -141,6 +141,17 @@ def test_main_shows_version():
             "multiple_fragments_client",
             CLIENTS_PATH / "multiple_fragments" / "expected_client",
         ),
+        (
+            (
+                CLIENTS_PATH / "shorter_results" / "pyproject.toml",
+                (
+                    CLIENTS_PATH / "shorter_results" / "queries.graphql",
+                    CLIENTS_PATH / "shorter_results" / "schema.graphql",
+                ),
+            ),
+            "shorter_results",
+            CLIENTS_PATH / "shorter_results" / "expected_client",
+        ),
     ],
     indirect=["project_dir"],
 )
diff --git a/tests/plugins/test_shorter_results.py b/tests/plugins/test_shorter_results.py
new file mode 100644
index 0000000..b18d510
--- /dev/null
+++ b/tests/plugins/test_shorter_results.py
@@ -0,0 +1,26 @@
+import ast
+
+from graphql import GraphQLSchema
+
+from ariadne_codegen.contrib.shorter_results import ShorterResultsPlugin
+
+
+def test_add_import():
+    module = ast.parse(
+        """
+from foo import bar
+
+class Client:
+    pass
+"""
+    )
+
+    plugin = ShorterResultsPlugin(GraphQLSchema(), {})
+    plugin.extended_imports = {
+        "foo": {"baz"},
+    }
+
+    updated = plugin.generate_client_module(module)
+    alias_names = [x.name for x in updated.body[0].names]
+
+    assert "baz" in alias_names