diff --git a/optype/numpy/_sequence_nd.py b/optype/numpy/_sequence_nd.py new file mode 100644 index 00000000..c00e1403 --- /dev/null +++ b/optype/numpy/_sequence_nd.py @@ -0,0 +1,12 @@ +from typing import Any, Protocol, TypeVar + + +T_co = TypeVar("T_co", covariant=True) + + +class SequenceND(Protocol[T_co]): + # a slimmed down version of `numpy._typing._NestedSequence` + def __len__(self, /) -> int: ... + def __getitem__(self, index: int, /) -> "T_co | SequenceND[T_co]": ... + def __contains__(self, x: object, /) -> bool: ... + def index(self, value: Any, /) -> int: ... # type: ignore[no-any-explicit] # pyright: ignore[reportAny, reportExplicitAny] diff --git a/optype/numpy/_to.py b/optype/numpy/_to.py index bb3c69ef..31a320a5 100644 --- a/optype/numpy/_to.py +++ b/optype/numpy/_to.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import sys from collections.abc import Sequence from typing import TypeAlias, TypeVar @@ -6,6 +8,7 @@ from ._array import CanArrayND from ._scalar import floating, integer, number +from ._sequence_nd import SequenceND if sys.version_info >= (3, 13): @@ -40,33 +43,37 @@ T = TypeVar("T") ST = TypeVar("ST", bound=np.generic) + _To1D: TypeAlias = CanArrayND[ST] | Sequence[T | ST] _To2D: TypeAlias = CanArrayND[ST] | Sequence[_To1D[ST, T]] +# recursive sequence type cannot be used due to a mypy bug: +# https://github.com/python/mypy/issues/18184 +_ToND: TypeAlias = CanArrayND[ST] | SequenceND[T | ST] | SequenceND[CanArrayND[ST]] ToScalar: TypeAlias = complex | bytes | str | np.generic ToArray1D: TypeAlias = _To1D[np.generic, complex | bytes | str] ToArray2D: TypeAlias = _To2D[np.generic, complex | bytes | str] -ToArrayND: TypeAlias = _To1D[np.generic, complex | bytes | str] | Sequence["ToArrayND"] +ToArrayND: TypeAlias = _ToND[np.generic, complex | bytes | str] ToBool: TypeAlias = bool | np.bool_ ToBool1D: TypeAlias = _To1D[np.bool_, bool] ToBool2D: TypeAlias = _To2D[np.bool_, bool] -ToBoolND: TypeAlias = _To1D[np.bool_, bool] | Sequence["ToBoolND"] +ToBoolND: TypeAlias = _ToND[np.bool_, bool] -_i_co = TypeAliasType("_i_co", integer | np.bool_) # type: ignore[no-any-explicit] -ToInt: TypeAlias = int | _i_co -ToInt1D: TypeAlias = _To1D[_i_co, int] -ToInt2D: TypeAlias = _To2D[_i_co, int] -ToIntND: TypeAlias = _To1D[_i_co, int] | Sequence["ToIntND"] +integer_co = TypeAliasType("integer_co", integer | np.bool_) # type: ignore[no-any-explicit] +ToInt: TypeAlias = int | integer_co +ToInt1D: TypeAlias = _To1D[integer_co, int] +ToInt2D: TypeAlias = _To2D[integer_co, int] +ToIntND: TypeAlias = _ToND[integer_co, int] -_f_co = TypeAliasType("_f_co", floating | integer | np.bool_) # type: ignore[no-any-explicit] -ToFloat: TypeAlias = float | _f_co -ToFloat1D: TypeAlias = _To1D[_f_co, float] -ToFloat2D: TypeAlias = _To2D[_f_co, float] -ToFloatND: TypeAlias = _To1D[_f_co, float] | Sequence["ToFloatND"] +floating_co = TypeAliasType("floating_co", floating | integer | np.bool_) # type: ignore[no-any-explicit] +ToFloat: TypeAlias = float | floating_co +ToFloat1D: TypeAlias = _To1D[floating_co, float] +ToFloat2D: TypeAlias = _To2D[floating_co, float] +ToFloatND: TypeAlias = _ToND[floating_co, float] -_c_co = TypeAliasType("_c_co", number | np.bool_) # type: ignore[no-any-explicit] -ToComplex: TypeAlias = complex | _c_co -ToComplex1D: TypeAlias = _To1D[_c_co, complex] -ToComplex2D: TypeAlias = _To2D[_c_co, complex] -ToComplexND: TypeAlias = _To1D[_c_co, complex] | Sequence["ToComplexND"] +complexfloating_co = TypeAliasType("complexfloating_co", number | np.bool_) # type: ignore[no-any-explicit] +ToComplex: TypeAlias = complex | complexfloating_co +ToComplex1D: TypeAlias = _To1D[complexfloating_co, complex] +ToComplex2D: TypeAlias = _To2D[complexfloating_co, complex] +ToComplexND: TypeAlias = _ToND[complexfloating_co, complex] diff --git a/tests/numpy/test_to.pyi b/tests/numpy/test_to.pyi index ed3cb072..b62d452d 100644 --- a/tests/numpy/test_to.pyi +++ b/tests/numpy/test_to.pyi @@ -4,9 +4,9 @@ import numpy as np import optype.numpy as onp vec_b: onp.Array1D[np.bool] | list[bool | np.bool_] -vec_i: onp.Array1D[np.intp] | list[int | np.intp] -vec_f: onp.Array1D[np.longdouble] | list[float | np.longdouble] -vec_c: onp.Array1D[np.clongdouble] | list[complex | np.clongdouble] +vec_i: onp.Array1D[np.intp] | list[int | np.intp] | range +vec_f: onp.Array1D[np.longdouble] | list[float | np.longdouble] | range +vec_c: onp.Array1D[np.clongdouble] | list[complex | np.clongdouble] | range mat_b: onp.Array2D[np.bool] | list[list[bool | np.bool_]] mat_i: onp.Array2D[np.intp] | list[list[int | np.intp]] @@ -29,26 +29,68 @@ si1: onp.ToInt = np.int_(42) sf0: onp.ToFloat = 42.0 sf1: onp.ToFloat = np.longdouble(42.0) -sc0: onp.ToComplex = 42.0 + 1j +sc0: onp.ToComplex = 42 + 0j sc1: onp.ToComplex = np.clongdouble(42.0 + 1j) # vectors - vb: onp.ToBool1D = vec_b vi: onp.ToInt1D = vec_i vf: onp.ToFloat1D = vec_f vc: onp.ToComplex1D = vec_c -# matrices +vsb: onp.ToBool1D = True # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +vsi: onp.ToInt1D = 42 # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +vsf: onp.ToFloat1D = 42.0 # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +vsc: onp.ToComplex1D = 42 + 1j # type: ignore[assignment] # pyright: ignore[reportAssignmentType] + +vub: onp.ToBool1D = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +vui: onp.ToInt1D = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +vuf: onp.ToFloat1D = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +vuc: onp.ToComplex1D = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +# matrices mb: onp.ToBool2D = mat_b mi: onp.ToInt2D = mat_i mf: onp.ToFloat2D = mat_f mc: onp.ToComplex2D = mat_c -# tensors +mvb: onp.ToBool2D = vec_b # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +mvi: onp.ToInt2D = vec_i # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +mvf: onp.ToFloat2D = vec_f # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +mvc: onp.ToComplex2D = vec_c # type: ignore[assignment] # pyright: ignore[reportAssignmentType] + +msb: onp.ToBool2D = True # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +msi: onp.ToInt2D = 42 # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +msf: onp.ToFloat2D = 42.0 # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +msc: onp.ToComplex2D = 42 + 0j # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +mub: onp.ToBool2D = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +mui: onp.ToInt2D = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +muf: onp.ToFloat2D = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +muc: onp.ToComplex2D = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType] + +# tensors tb: onp.ToBoolND = arr_b ti: onp.ToIntND = arr_i tf: onp.ToFloatND = arr_f tc: onp.ToComplexND = arr_c + +tmb: onp.ToBoolND = mat_b +tmi: onp.ToIntND = mat_i +tmf: onp.ToFloatND = mat_f +tmc: onp.ToComplexND = mat_c + +tvb: onp.ToBoolND = vec_b +tvi: onp.ToIntND = vec_i +tvf: onp.ToFloatND = vec_f +tvc: onp.ToComplexND = vec_c + +tsb: onp.ToBoolND = True # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +tsi: onp.ToIntND = 42 # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +tsf: onp.ToFloatND = 42.0 # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +tsc: onp.ToComplexND = 42 + 0j # type: ignore[assignment] # pyright: ignore[reportAssignmentType] + +tub: onp.ToBoolND = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +tui: onp.ToIntND = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +tuf: onp.ToFloatND = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType] +tuc: onp.ToComplexND = "illegal" # type: ignore[assignment] # pyright: ignore[reportAssignmentType]