Mypy fails with Expected type 'Type[T]', got 'UnionType' instead

I'm using pydantic to load data from JSON files that can have different parameters, for this example either x or y. I use a Loader class to load them into pydantic models so that I can validate them at runtime but I cannot make my code compliant with mypy. It fails with this error: Expected type 'Type[T]', got 'UnionType' instead.

I have mypy 1.10.0 and pydantic 2.7.0.

Here is a reproducible example of my use case:

from typing import Iterator, Type, TypeVar, Union

from pydantic import BaseModel, TypeAdapter


class A(BaseModel):
    x: int


class B(BaseModel):
    y: int


T = TypeVar("T", bound=BaseModel)


class Loader:
    def load_models(self, model_class: Type[T]) -> Iterator[T]:
        # logic to load models
        my_dict = {"x": 1}  # fake data for testing
        yield TypeAdapter(model_class).validate_python(my_dict)
        my_dict = {"y": 2}  # fake data for testing
        yield TypeAdapter(model_class).validate_python(my_dict)


ModelAOrB = A | B

print(list(Loader().load_models(ModelAOrB)))

Is there a way to make it work with mypy?

I found a way to do what I wanted thanks to this post, the key was to use RootModel from pydantic and use it as an argument directly. Without forgetting to return .root, mypy now accepts the following code:

class Loader:
    def load_models(self, root_model: type[RootModel[T]]) -> Iterator[T]:
        # logic to load models
        my_dict = {"x": 1}  # fake data for testing
        yield root_model.model_validate(my_dict).root # or root_model(**my_dict).root
        my_dict = {"y": 2}  # fake data for testing
        yield root_model.model_validate(my_dict).root # or root_model(**my_dict).root


ModelAOrB = A | B

print(list(Loader().load_models(RootModel[ModelAOrB])))