reified_generics.py

from __future__ import annotations

from typing import (
    _GenericAlias, Annotated, Any, Generic, Protocol, TypeVar,  # type: ignore[attr-defined]
)

T = TypeVar("T")


class _ReifiedGenericAlias(_GenericAlias, _root=True):
    def __subclasscheck__(self, subclass: Annotated[Any, "could be any random class"] | _ReifiedGenericAlias) -> bool:
        # could be any random class, check it first
        if not issubclass(subclass.__origin__, self.__origin__):
            return False
        for parameter, self_arg, subclass_arg in zip(self.__origin__.__parameters__,
                self.__args__, subclass.__args__):
            if parameter.__covariant__ and not issubclass(subclass_arg, self_arg):
                return False
            elif parameter.__contravariant__ and not issubclass(self_arg, subclass_arg):
                return False
            elif self_arg is not subclass_arg:
                return False
        return True

    def __instancecheck__(self, instance: ReifiedGeneric | Annotated[Any, "could be any random instance"]) -> bool:
        # could be any random instance, check it first
        if not isinstance(instance, self.__origin__):
            return False
        for parameter, self_arg, subclass_arg in zip(self.__origin__.__parameters__,
                self.__args__, instance.__orig_class__.__args__):
            if parameter.__covariant__ and not issubclass(subclass_arg, self_arg):
                return False
            elif parameter.__contravariant__ and not issubclass(self_arg, subclass_arg):
                return False
            elif self_arg is not subclass_arg:
                return False
        return True


class OrigClass(Protocol):
    __args__: tuple[type]
    __parameters__: tuple[TypeVar]


class ReifiedGeneric(Generic[T]):
    """
    TODO: fail if generics not provided
    """
    def __init__(self):
        self.__orig_class__: OrigClass

    def __class_getitem__(cls, item) -> type[ReifiedGeneric[T]]:
        result = super().__class_getitem__(item)
        return _ReifiedGenericAlias(cls, result.__args__)


# TESTING ------------------------------------------------------------------------------

T2 = TypeVar("T2")



class RA(ReifiedGeneric[tuple[T, T2]]):
    ...


class RList(list[T], ReifiedGeneric[tuple[T]]):
    ...


class A(Generic[T, T2]):
    ...


def test_args_and_params():
    assert A[T, T2].__args__ == RA[T, T2].__args__
    assert A[T, T2].__parameters__ == RA[T, T2].__parameters__
    assert A[int, T2].__args__ == RA[int, T2].__args__
    assert A[int, T2].__parameters__ == RA[int, T2].__parameters__


def test_reified_list():
    it = RList[int]([1, 2, 3]).__orig_class__
    assert it.__origin__ == list
    assert it.__args__ == (int,)
    assert it.__parameters__ == ()

def test_isinstance():
    assert isinstance(RA[int, str](), RA[int, str])
    assert not isinstance(RA[int, str](), RA[int, int])

def test_issubclass():
    assert isinstance(RA[int, str], RA[int, str])
    assert not isinstance(RA[int, str], RA[int, int])