from __future__ import annotations
|
|
import sys
|
import types
|
from collections.abc import Generator, Iterable, Iterator
|
from typing import (
|
Any,
|
ClassVar,
|
NoReturn,
|
TypeVar,
|
TYPE_CHECKING,
|
)
|
|
import numpy as np
|
|
__all__ = ["_GenericAlias", "NDArray"]
|
|
_T = TypeVar("_T", bound="_GenericAlias")
|
|
|
def _to_str(obj: object) -> str:
|
"""Helper function for `_GenericAlias.__repr__`."""
|
if obj is Ellipsis:
|
return '...'
|
elif isinstance(obj, type) and not isinstance(obj, _GENERIC_ALIAS_TYPE):
|
if obj.__module__ == 'builtins':
|
return obj.__qualname__
|
else:
|
return f'{obj.__module__}.{obj.__qualname__}'
|
else:
|
return repr(obj)
|
|
|
def _parse_parameters(args: Iterable[Any]) -> Generator[TypeVar, None, None]:
|
"""Search for all typevars and typevar-containing objects in `args`.
|
|
Helper function for `_GenericAlias.__init__`.
|
|
"""
|
for i in args:
|
if hasattr(i, "__parameters__"):
|
yield from i.__parameters__
|
elif isinstance(i, TypeVar):
|
yield i
|
|
|
def _reconstruct_alias(alias: _T, parameters: Iterator[TypeVar]) -> _T:
|
"""Recursively replace all typevars with those from `parameters`.
|
|
Helper function for `_GenericAlias.__getitem__`.
|
|
"""
|
args = []
|
for i in alias.__args__:
|
if isinstance(i, TypeVar):
|
value: Any = next(parameters)
|
elif isinstance(i, _GenericAlias):
|
value = _reconstruct_alias(i, parameters)
|
elif hasattr(i, "__parameters__"):
|
prm_tup = tuple(next(parameters) for _ in i.__parameters__)
|
value = i[prm_tup]
|
else:
|
value = i
|
args.append(value)
|
|
cls = type(alias)
|
return cls(alias.__origin__, tuple(args), alias.__unpacked__)
|
|
|
class _GenericAlias:
|
"""A python-based backport of the `types.GenericAlias` class.
|
|
E.g. for ``t = list[int]``, ``t.__origin__`` is ``list`` and
|
``t.__args__`` is ``(int,)``.
|
|
See Also
|
--------
|
:pep:`585`
|
The PEP responsible for introducing `types.GenericAlias`.
|
|
"""
|
|
__slots__ = (
|
"__weakref__",
|
"_origin",
|
"_args",
|
"_parameters",
|
"_hash",
|
"_starred",
|
)
|
|
@property
|
def __origin__(self) -> type:
|
return super().__getattribute__("_origin")
|
|
@property
|
def __args__(self) -> tuple[object, ...]:
|
return super().__getattribute__("_args")
|
|
@property
|
def __parameters__(self) -> tuple[TypeVar, ...]:
|
"""Type variables in the ``GenericAlias``."""
|
return super().__getattribute__("_parameters")
|
|
@property
|
def __unpacked__(self) -> bool:
|
return super().__getattribute__("_starred")
|
|
@property
|
def __typing_unpacked_tuple_args__(self) -> tuple[object, ...] | None:
|
# NOTE: This should return `__args__` if `__origin__` is a tuple,
|
# which should never be the case with how `_GenericAlias` is used
|
# within numpy
|
return None
|
|
def __init__(
|
self,
|
origin: type,
|
args: object | tuple[object, ...],
|
starred: bool = False,
|
) -> None:
|
self._origin = origin
|
self._args = args if isinstance(args, tuple) else (args,)
|
self._parameters = tuple(_parse_parameters(self.__args__))
|
self._starred = starred
|
|
@property
|
def __call__(self) -> type[Any]:
|
return self.__origin__
|
|
def __reduce__(self: _T) -> tuple[
|
type[_T],
|
tuple[type[Any], tuple[object, ...], bool],
|
]:
|
cls = type(self)
|
return cls, (self.__origin__, self.__args__, self.__unpacked__)
|
|
def __mro_entries__(self, bases: Iterable[object]) -> tuple[type[Any]]:
|
return (self.__origin__,)
|
|
def __dir__(self) -> list[str]:
|
"""Implement ``dir(self)``."""
|
cls = type(self)
|
dir_origin = set(dir(self.__origin__))
|
return sorted(cls._ATTR_EXCEPTIONS | dir_origin)
|
|
def __hash__(self) -> int:
|
"""Return ``hash(self)``."""
|
# Attempt to use the cached hash
|
try:
|
return super().__getattribute__("_hash")
|
except AttributeError:
|
self._hash: int = (
|
hash(self.__origin__) ^
|
hash(self.__args__) ^
|
hash(self.__unpacked__)
|
)
|
return super().__getattribute__("_hash")
|
|
def __instancecheck__(self, obj: object) -> NoReturn:
|
"""Check if an `obj` is an instance."""
|
raise TypeError("isinstance() argument 2 cannot be a "
|
"parameterized generic")
|
|
def __subclasscheck__(self, cls: type) -> NoReturn:
|
"""Check if a `cls` is a subclass."""
|
raise TypeError("issubclass() argument 2 cannot be a "
|
"parameterized generic")
|
|
def __repr__(self) -> str:
|
"""Return ``repr(self)``."""
|
args = ", ".join(_to_str(i) for i in self.__args__)
|
origin = _to_str(self.__origin__)
|
prefix = "*" if self.__unpacked__ else ""
|
return f"{prefix}{origin}[{args}]"
|
|
def __getitem__(self: _T, key: object | tuple[object, ...]) -> _T:
|
"""Return ``self[key]``."""
|
key_tup = key if isinstance(key, tuple) else (key,)
|
|
if len(self.__parameters__) == 0:
|
raise TypeError(f"There are no type variables left in {self}")
|
elif len(key_tup) > len(self.__parameters__):
|
raise TypeError(f"Too many arguments for {self}")
|
elif len(key_tup) < len(self.__parameters__):
|
raise TypeError(f"Too few arguments for {self}")
|
|
key_iter = iter(key_tup)
|
return _reconstruct_alias(self, key_iter)
|
|
def __eq__(self, value: object) -> bool:
|
"""Return ``self == value``."""
|
if not isinstance(value, _GENERIC_ALIAS_TYPE):
|
return NotImplemented
|
return (
|
self.__origin__ == value.__origin__ and
|
self.__args__ == value.__args__ and
|
self.__unpacked__ == getattr(
|
value, "__unpacked__", self.__unpacked__
|
)
|
)
|
|
def __iter__(self: _T) -> Generator[_T, None, None]:
|
"""Return ``iter(self)``."""
|
cls = type(self)
|
yield cls(self.__origin__, self.__args__, True)
|
|
_ATTR_EXCEPTIONS: ClassVar[frozenset[str]] = frozenset({
|
"__origin__",
|
"__args__",
|
"__parameters__",
|
"__mro_entries__",
|
"__reduce__",
|
"__reduce_ex__",
|
"__copy__",
|
"__deepcopy__",
|
"__unpacked__",
|
"__typing_unpacked_tuple_args__",
|
"__class__",
|
})
|
|
def __getattribute__(self, name: str) -> Any:
|
"""Return ``getattr(self, name)``."""
|
# Pull the attribute from `__origin__` unless its
|
# name is in `_ATTR_EXCEPTIONS`
|
cls = type(self)
|
if name in cls._ATTR_EXCEPTIONS:
|
return super().__getattribute__(name)
|
return getattr(self.__origin__, name)
|
|
|
# See `_GenericAlias.__eq__`
|
if sys.version_info >= (3, 9):
|
_GENERIC_ALIAS_TYPE = (_GenericAlias, types.GenericAlias)
|
else:
|
_GENERIC_ALIAS_TYPE = (_GenericAlias,)
|
|
ScalarType = TypeVar("ScalarType", bound=np.generic, covariant=True)
|
|
if TYPE_CHECKING or sys.version_info >= (3, 9):
|
_DType = np.dtype[ScalarType]
|
NDArray = np.ndarray[Any, np.dtype[ScalarType]]
|
else:
|
_DType = _GenericAlias(np.dtype, (ScalarType,))
|
NDArray = _GenericAlias(np.ndarray, (Any, _DType))
|
