Graphene input error for Pydantic models with discriminator while generating Input object schema

Question

I am using pydantic validations for my requirements and it uses discriminator. I am writing GraphQL APIs and want to convert those pydantic models into graphene input objects.

Below is my code.

from graphene_pydantic import PydanticInputObjectType, PydanticObjectType
import graphene
from typing import Literal, Union
from pydantic import BaseModel, Field


class Cat(BaseModel):
    pet_type: Literal['cat']
    meows: int


class Dog(BaseModel):
    pet_type: Literal['dog']
    barks: float


class Lizard(BaseModel):
    pet_type: Literal['reptile', 'lizard']
    scales: bool


class Model(BaseModel):
    pet: Union[Cat, Dog, Lizard] = Field(..., discriminator='pet_type')
    n: int


print(Model(pet={'pet_type': 'dog', 'barks': 3.14, 'eats': 'biscuit'}, n=1))


class Input(PydanticInputObjectType):
    class Meta:
        model = Model
        # exclude specified fields
        exclude_fields = ("id",)


class Output(PydanticObjectType):
    class Meta:
        model = Model
        # exclude specified fields
        exclude_fields = ("id",)


class CreateAnimal(graphene.Mutation):
    class Arguments:
        input = Input()

    output = Output

    @staticmethod
    def mutate(parent, info, input):
        print(input)
        # save model here
        return input


class Mutation(graphene.ObjectType):
    createPerson = CreateAnimal.Field()


schema = graphene.Schema(mutation=Mutation)
print(schema)

I tried by commenting on the discriminator code and it's working fine but I need those validations for graphql also. If I run that code it's throwing the below error.

File "\AppData\Local\Programs\Python\Python310\lib\site-packages\graphql\type\definition.py", line 1338, in fields    raise TypeError(f"{self.name} fields cannot be resolved. {error}")
TypeError: Input fields cannot be resolved. The input field type must be a GraphQL input type.

Can someone help me with this?

I am using graphene-pydantic for this.

Answer 1

Hey i havent tested your code but i think this what you need:

class CreateAnimal(graphene.Mutation):
    class Arguments:
        input = graphene.Argument(Input)

    output = graphene.Field(Output)