ENH: Replace `NestedSequence` with a proper nested sequence protocol
@@ -9,6 +9,8 @@ valid for inputs that match the given type annotations.
from cupy.cuda import Device as _Device
+from __future__ import annotations
+
__all__ = [
"Array",
"Device",
@@ -19,7 +21,16 @@ __all__ = [
]
import sys
-from typing import Any, Literal, Sequence, Type, Union, TYPE_CHECKING, TypeVar
+from typing import (
+ Any,
+ Literal,
+ Sequence,
+ Type,
+ Union,
+ TYPE_CHECKING,
+ TypeVar,
+ Protocol,
+)
from ._array_object import Array
from numpy import (
@@ -36,10 +47,11 @@ from numpy import (
float64,
)
-# This should really be recursive, but that isn't supported yet. See the
-# similar comment in numpy/typing/_array_like.py
-_T = TypeVar("_T")
-NestedSequence = Sequence[Sequence[_T]]
+_T_co = TypeVar("_T_co", covariant=True)
+
+class NestedSequence(Protocol[_T_co]):
+ def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
+ def __len__(self, /) -> int: ...
Device = _Device
if TYPE_CHECKING or sys.version_info >= (3, 9):
