GraphQL @OneOf directive implementation in Python using graphql-core

one_of.py

"""
This example shows an implementation of the @oneOf directive proposed here:
https://github.com/graphql/graphql-spec/pull/825.

It requires Python 3.9 and graphql-core 3.1.3 (https://github.com/graphql-python/graphql-core).

Although the directive support can be implemented out of the box when used in a field
definition, its use in input object type definition requires a specific graphql-core
feature custom output types of input object types
(https://graphql-core-3.readthedocs.io/en/latest/diffs.html#custom-output-types-of-input-object-types).

This feature allows converting an input object (parsed as a dictionary) into a custom class
instance (because using objects and attributes is more pythonic than using dictionaries,
contrary to Javascript where objects and dictionaries are pretty much the same things);
but this feature can also be used to validate an input object.

However, there is an issue still opened in graphql-js repository:
https://github.com/graphql/graphql-js/issues/361, which proposes to add validation to
GraphQLInputObjectType. If this feature were released, it would allows to support the
directive directly with the reference implementation.

The directive support is added to the schema through a visit of it, which will modify the
decorated nodes:
- for field definition, the resolver is altered by checking that only one argument is
  provided before executing it;
- for input object type definition, validation function (see above) is added to check
  that only one field is provided.
"""
from typing import Any, Callable, Collection, Optional, TypeVar, cast

from graphql import (
    DirectiveNode,
    ExecutionResult,
    GraphQLInputObjectType,
    GraphQLObjectType,
    GraphQLSchema,
    build_schema,
    default_field_resolver,
    graphql_sync,
)

#########
# Utils #
#########

Resolve = TypeVar("Resolve")  # Please, forgive Python typing verbosity.


def bind(schema: GraphQLSchema, type: str, field: str) -> Callable[[Resolve], Resolve]:
    """Bind a resolver to a field in the schema."""

    def decorator(resolve: Resolve) -> Resolve:
        cast(GraphQLObjectType, schema.type_map[type]).fields[field].resolve = resolve
        return resolve

    return decorator


def check_one_of(data: dict[str, Any]) -> dict[str, Any]:
    """Simply check that a dictionary contains one and only one key."""
    if len(data) != 1:
        raise ValueError(f"@oneOf expects only one key, found {list(data)}")
    return data


def contains_one_of_directive(directives: Optional[Collection[DirectiveNode]]) -> bool:
    """Return if directives contains @oneOf"""
    return directives and any(d.name.value == "oneOf" for d in directives)


def handle_one_of(schema: GraphQLSchema):
    """Alter the schema in order to support @oneOf directive."""
    # Iterate on the schema types, looking for decorated GraphQLInputObjectType and
    # GraphQLObjectType with decorated field.
    for type_ in schema.type_map.values():
        if isinstance(type_, GraphQLInputObjectType):
            # Alter the type validator
            if contains_one_of_directive(type_.ast_node.directives):
                # If type.out_type is not the default one, compose it with check_one_of,
                # otherwise, just replace it.
                if type_.out_type is not GraphQLInputObjectType.out_type:
                    wrapped = type_.out_type
                    type_.out_type = lambda data: wrapped(check_one_of(data))
                else:
                    type_.out_type = check_one_of
        if isinstance(type_, GraphQLObjectType):
            for field in type_.fields.values():
                # /!\ Introspection fields doesn't have ast_node
                if field.ast_node is not None and contains_one_of_directive(
                    field.ast_node.directives
                ):
                    # Alter the field resolver
                    resolve = field.resolve or default_field_resolver
                    field.resolve = lambda obj, info, /, **kwargs: resolve(
                        obj, info, **check_one_of(kwargs)
                    )


def format_execution_result(execution_result: ExecutionResult) -> dict[str, Any]:
    """Format execution result in order to be verifiable in tests"""
    return {
        "data": execution_result.data,
        "errors": (
            [str(err) for err in execution_result.errors]
            if execution_result.errors is not None
            else None
        ),
    }


###################################
# Schema declaration and bindings #
###################################

schema_sdl = """
directive @oneOf on FIELD_DEFINITION | INPUT_OBJECT

scalar User

type Query {
  user(
    id: ID
    email: String
    username: String
    registrationNumber: Int
  ): User @oneOf
}

input PetInput @oneOf {
  cat: CatInput
  dog: DogInput
  fish: FishInput
}

input CatInput { name: String!, numberOfLives: Int }
input DogInput { name: String!, wagsTail: Boolean }
input FishInput { name: String!, bodyLengthInMm: Int }

scalar Pet

type Mutation {
  addPet(pet: PetInput!): Pet
}
"""

schema = build_schema(schema_sdl)


@bind(schema, "Query", "user")
def resolve_user(_, info, **kwargs):
    assert len(kwargs) == 1
    return kwargs


@bind(schema, "Mutation", "addPet")
def resolve_add_pet(_, info, pet: dict[str, Any]):
    assert len(pet) == 1
    return pet


handle_one_of(schema)


#########
# Tests #
#########

def execute(operation: str) -> dict[str, Any]:
    """Execute an operation and format its result"""
    return format_execution_result(graphql_sync(schema, operation))


# Test Query.user with 1 argument, no error
assert execute("query {user(registrationNumber: 42)}") == {
    "data": {"user": {"registrationNumber": 42}},
    "errors": None,
}
# Test Query.user with 2 arguments, error!
assert execute('query {user(username: "wyfo", registrationNumber: 42)}') == {
    "data": {"user": None},
    "errors": [
        """\
@oneOf expects only one key, found ['username', 'registrationNumber']

GraphQL request:1:8
1 | query {user(username: \"wyfo\", registrationNumber: 42)}
  |        ^"""
    ],
}

# Test Mutation.addPet with 1 PetInput fields, no error
assert execute('mutation {addPet(pet: {cat: {name: "wyfo"}})}') == {
    "data": {"addPet": {"cat": {"name": "wyfo"}}},
    "errors": None,
}
# Test Mutation.addPet with 2 PetInput fields, error!
assert execute(
    'mutation {addPet(pet: {cat: {name: "wyfo"}, dog: {name: "wyfo"}})}'
) == {
    "data": {"addPet": None},
    "errors": [
        """\
@oneOf expects only one key, found ['cat', 'dog']

GraphQL request:1:11
1 | mutation {addPet(pet: {cat: {name: \"wyfo\"}, dog: {name: \"wyfo\"}})}
  |           ^"""
    ],
}