What are valid things to return from __class_getitem__?

[adriangb]

What are valid things to return from __class_getitem__? PEP 560 gives as an example a string of the form f"{cls.__name__}[{item.__name__}]". Subclasses of typing.Generic return a typing._GenericAlias instance. I tried to write a class that tracks generic parameters so that they are available when it is instantiated by returning itself (code below) and am running into problems.

The issue I've encountered is that the internals of typing._GenericAlias look for type vars in their arguments using typing._collect_type_vars. So List[Dict[K, V]] -> [~K, ~V] because List is a _GenericAlias and thus _GenericAlias.__class_getitem__ -> _GenericAlias.__init__ -> typing._collect_type_vars get called). In both typing.py and typing_extensions _collect_type_vars specifically checks if each type is an instance of _GenericAlias and only collects type vars from it (by accessing __parameters__) if it is (typing_extensions implementation). So for my class, which has or can have __parameters__ but is not a subclass of _GenericAlias no type variables are collected. This eventually triggers this error if you try to do something like List[Model[T]][int] because the outer List thinks it has no generic parameters.

Would it be unreasonable for _GenericAlias to do getattr(type, '__parameters__', ()) or something like that? Or is there some way I can restructure my implementation to avoid these issues?

from typing import Any, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
from typing_extensions import dataclass_transform

TypeVarMap = Dict[Type[Any], Type[Any]]

@dataclass_transform()
class BaseModel:
    __type_var_map__: TypeVarMap = {}
    __origin__: Optional[Type[Any]] = None

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        print(self.__type_var_map__)

    def __init_subclass__(cls) -> None:
        super().__init_subclass__()
        cls.__origin__ = cls
        if issubclass(cls, Generic):
            cls.__type_var_map__ = cls.__type_var_map__ or {
                tv: tv
                for tv in getattr(cls, '__parameters__', {})
            }

    def __class_getitem__(cls, args: Union[Type[Any], Tuple[Type[Any], ...]]) -> Any:
        if not isinstance(args, tuple):
            tvs = [args]
        else:
            tvs = list(args)
        tv_map: TypeVarMap = cls.__type_var_map__.copy()
        missing: TypeVarMap = {k: v for k, v in tv_map.items() if isinstance(v, TypeVar)}
        tv_map = {**tv_map, **{k: new_v for (k, _), new_v in zip(missing.items(), tvs)}}
        return type(cls.__name__, cls.__bases__, {**cls.__dict__, '__type_var_map__': tv_map})

T = TypeVar('T')

class Model(BaseModel, Generic[T]):
    x: T

partial_model = Model[T]
concrete_model1 = partial_model[str]
print(concrete_model1.__type_var_map__)
concrete_model2 = partial_model[int]
print(concrete_model2.__type_var_map__)

partial_model_list = List[Model[T]]
partial_model_list[int]  # error

This is somewhat similar (in goal at least) to #1050 but I didn't want to hijack that.

[dmontagu]

I want to provide some additional context to @adriangb's question.

When creating a web api, it's common to have many types that have the same top-level fields, with a reused inner field. For example:

from typing import Generic, List, Optional, TypeVar
from pydantic import BaseModel


class IntResponse(BaseModel):
    value: Optional[int]
    error_message: Optional[str] = None


class StrResponse(BaseModel):
    value: Optional[str]
    error_message: Optional[str] = None

# You could imagine other kinds of XyzResponse classes

You can clearly see this looks ripe for some sort of generics-like approach. However, one important difference with the way generics work in the python typing module is that when you parametrize a typing-generics class, you don't retain the ability to have different underlying behavior at runtime as a function of the parameters. In particular, for the above (non-generic) classes, we have:

IntResponse.model_validate({'value': 1})  # passes
StrResponse.model_validate({'value': 'a'})  # passes
# IntResponse.model_validate({'value': 'a'})  # would raise an error

As far as I can tell, because of the way typing.Generic works, without redefining __class_getitem__ to return something other than a _GenericAlias or heavily abusing typing-internal APIs related to _GenericAlias, there is nothing I could insert into the StandardGenericResponse class body below that would make the following code produce the same results as we get in the code block above:

T = TypeVar("T")
class StandardGenericResponse(Generic[T]):
    value: Optional[T]
    error_message: Optional[str] = None

    @classmethod
    def model_validate(cls, data):
        ...  # ???

    ...  # any other methods/etc. as appropriate

StandardGenericResponse[int].model_validate({'value': 1})  # *should not* error at runtime
StandardGenericResponse[str].model_validate({'value': 'a'})  # *should not* error at runtime
StandardGenericResponse[int].model_validate({'value': 'a'})  # *should* raise an error at runtime

However, through some elbow grease, we have largely been able to achieve a work-around for this in pydantic, by adding a bunch of custom __class_getitem__ machinery, which just so happens to play nice with typecheckers and IDEs:

T = TypeVar("T")

class ResponseModel(BaseModel, Generic[T]):  # BaseModel has a custom implementation of __class_getitem__
    value: Optional[T]
    error_message: Optional[str] = None

ResponseModel[int].model_validate({'value': 1})  # passes
ResponseModel[str].model_validate({'value': 'a'})  # passes
# ResponseModel[int].model_validate({'value': 'a'})  # would raise an error

However, there are two issues we've run into with this:

• First, we've essentially had to copy all of the typing module's TypeVar substitution logic but with modifications to ensure that our "custom" generic classes are introspected properly. We were able to achieve this, but it would simplify things immensely if things like typing.get_args, typing.get_origin, etc., worked out of the box with our types.
  • (You can see our duplication of the type substitution logic in the pydantic repository in the vicinity of the function here.)
  • I'll note that we've used __pydantic_generic_origin__, __pydantic_generic_args__, etc. for clarity that it's our custom implementation, but could just as easily change back to using __origin__, __args__, etc., if that meant things would work better.
• More problematic — when typing types are at the "root" of the generic model, certain things that you'd expect to work stop working, largely because various functions in the typing module (such as get_args, get_origin, etc.) only read args, origin, parameters, etc. off classes of specific types private to the typing module. As an example:

ResponseModelList = List[ResponseModel[T]]
#   The following line would raise TypeError: typing.List[__main__.XModel] is not a generic class
# IntResponseModelList = ResponseModelList[int] 
#   Ideally, it would produce List[ResponseModel[int]]

List[StandardGenericResponse[T]][int]  # no error, produces List[StandardGenericResponse[int]]

---

With the above description as motivation, ultimately, I think the thing that would empower us to resolve all of these issues (and simplify our codebase significantly) would be for _GenericAlias — a private type in the typing module required as the input type for proper interaction with various typing utility functions — to be (functionally) replaced with a protocol that could be implemented by developers in their own classes and returned from __class_getitem__ for proper interaction with the python typing ecosystem.

In particular, if the following code could run without an error on the final line, I think we'd be in good shape:

from dataclasses import dataclass
from typing import Dict, List, Generic, TypeVar, Any

T = TypeVar("T")


@dataclass
class MyGenericAlias:
    __origin__: Any
    __parameters__: Any
    __args__: Any

    @staticmethod
    def construct(cls):
        return MyGenericAlias(cls.__origin__, cls.__parameters__, cls.__args__)

    def __call__(self, *args, **kwargs):
        # without this, we'll get:
        # TypeError: Parameters to generic types must be types. Got MyGenericAlias(__origin__=<class '__main__.CustomGeneric'>, __parameters__=(~T,), __args__=(typing.D.
        pass

class StandardGeneric(Generic[T]):
    pass


class CustomGeneric(Generic[T]):
    __origin__: Any
    __parameters__: Any
    __args__: Any

    def __class_getitem__(cls, item) -> MyGenericAlias:
        # return super().__class_getitem__(item)
        # Only difference: rather than
        return MyGenericAlias(
            __origin__=CustomGeneric,
            __parameters__=getattr(item, '__parameters__', None),
            __args__=(item,)
        )


print(MyGenericAlias.construct(StandardGeneric[Dict[int, T]]))
# the following looks basically the same, indicating same origin, parameters, args
print(MyGenericAlias.construct(CustomGeneric[Dict[int, T]]))

List[StandardGeneric[Dict[int, T]]][int]
# The next line raises:
# TypeError: typing.List[MyGenericAlias(__origin__=<class '__main__.CustomGeneric'>, __parameters__=(~T,), __args__=(typing.Dict[int, ~T],))] is not a generic class
List[CustomGeneric[Dict[int, T]]][int]

(Happy to explain more if useful.)