10000 ♻️ variadic associated rank types by jorenham · Pull Request #559 · numpy/numtype · GitHub
[go: up one dir, main page]
Skip to content

♻️ variadic associated rank types #559

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
May 19, 2025
Merged

♻️ variadic associated rank types #559

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
0a1d1fb
♻️ prefer associated-ish over generic types rank types
jorenham May 18, 2025
2b68d40
♻️ NEP 50 type alias shape type defaults to `Any`
jorenham May 18, 2025
01c603d
🐴 work around a mypy crash
jorenham May 18, 2025
5318fdd
✨ `_numtype.AnyShape`
jorenham May 18, 2025
103545a
♻️ variadic rank types
jorenham May 19, 2025
f95b93b
🤖 prefer `Shape` over `Rank` for unknown-sized array types
jorenham May 19, 2025
c2bee24
🐴 work around at least 5 mypy bugs (not even kidding)
jorenham May 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/_numtype/@test/generated/test_rank.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# @generated 2025-05-14T01:47:22Z with tool/testgen.py
# @generated 2025-05-19T03:25:34Z with tool/testgen.py
from typing import Any

import _numtype as _nt
Expand Down Expand Up @@ -31,7 +31,7 @@ r0n_le_r0n: _nt.HasRankLE[_nt.Rank0N] = r0n

r0_ge_s0: _nt.HasRankGE[_nt.Shape0] = r0
r0_ge_r0: _nt.HasRankGE[_nt.Rank0] = r0
r0_ge_s0n: _nt.HasRankGE[_nt.Shape0N] = r0
r0_ge_s0n: _nt.HasRankGE[_nt.Shape0N] = r0 # type: ignore[assignment]
r0_ge_r0n: _nt.HasRankGE[_nt.Rank0N] = r0
r0n_ge_s0: _nt.HasRankGE[_nt.Shape0] = r0n
r0n_ge_r0: _nt.HasRankGE[_nt.Rank0] = r0n
Expand Down
25 changes: 25 additions & 0 deletions src/_numtype/@test/test_rank_shape.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
from typing import assert_type

import _numtype as _nt

# TODO: remove the `# type: ignore`s once python/mypy#19110 is fixed

a0: _nt.Array0D
assert_type(a0.__inner_shape__, _nt.Rank0)
assert_type(a0.shape, _nt.Shape0) # type: ignore[assert-type]

a1: _nt.Array1D
assert_type(a1.__inner_shape__, _nt.Rank1)
assert_type(a1.shape, _nt.Shape1) # type: ignore[assert-type]

a2: _nt.Array2D
assert_type(a2.__inner_shape__, _nt.Rank2)
assert_type(a2.shape, _nt.Shape2) # type: ignore[assert-type]

a3: _nt.Array3D
assert_type(a3.__inner_shape__, _nt.Rank3)
assert_type(a3.shape, _nt.Shape3) # type: ignore[assert-type]

a4: _nt.Array4D
assert_type(a4.__inner_shape__, _nt.Rank4)
assert_type(a4.shape, _nt.Shape4) # type: ignore[assert-type]
11 changes: 7 additions & 4 deletions src/_numtype/__init__.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,11 @@ from ._array import (
Array2D as Array2D,
Array3D as Array3D,
Array4D as Array4D,
ArrayND as ArrayND,
MArray as MArray,
MArray0D as MArray0D,
MArray1D as MArray1D,
MArray2D as MArray2D,
MArray3D as MArray3D,
MArrayND as MArrayND,
Matrix as Matrix,
StringArray as StringArray,
StringArray0D as StringArray0D,
Expand Down Expand Up @@ -86,6 +84,7 @@ from ._nep50 import (
CastsWithScalar as CastsWithScalar,
)
from ._rank import (
HasInnerShape as HasInnerShape,
HasRankGE as HasRankGE,
HasRankLE as HasRankLE,
Rank as Rank,
Expand Down Expand Up @@ -141,6 +140,7 @@ from ._scalar_co import (
co_ulong as co_ulong,
)
from ._shape import (
AnyShape as AnyShape,
Shape as Shape,
Shape0 as Shape0,
Shape0N as Shape0N,
Expand Down Expand Up @@ -170,7 +170,8 @@ _ToT = TypeVar("_ToT")

@type_check_only
class CanArray0D(Protocol[_ScalarT_co]):
def __array__(self, /) -> np.ndarray[Shape0, np.dtype[_ScalarT_co]]: ...
# TODO: remove `| Rank0` once python/mypy#19110 is fixed
def __array__(self, /) -> np.ndarray[Shape0 | Rank0, np.dtype[_ScalarT_co]]: ...

@type_check_only
class CanArray1D(Protocol[_ScalarT_co]):
Expand All @@ -186,11 +187,13 @@ class CanArray3D(Protocol[_ScalarT_co]):

@type_check_only
class CanArrayND(Protocol[_ScalarT_co]):
def __array__(self, /) -> np.ndarray[Shape, np.dtype[_ScalarT_co]]: ...
# TODO: remove `| Rank0` once python/mypy#19110 is fixed
def __array__(self, /) -> np.ndarray[Shape | Rank0, np.dtype[_ScalarT_co]]: ...

@type_check_only
class CanLenArrayND(Protocol[_ScalarT_co]):
def __len__(self, /) -> int: ...
# TODO: remove `| Rank0` once python/mypy#19110 is fixed
def __array__(self, /) -> np.ndarray[Shape, np.dtype[_ScalarT_co]]: ...

@type_check_only
Expand Down
13 changes: 5 additions & 8 deletions src/_numtype/_array.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ from typing_extensions import TypeAliasType, TypeVar

import numpy as np

from ._rank import Rank, Rank0, Rank1, Rank2, Rank3, Rank4
from ._shape import Shape
from ._rank import Rank0, Rank1, Rank2, Rank3, Rank4
from ._shape import AnyShape, Shape

__all__ = [
"Array",
Expand All @@ -15,13 +15,11 @@ __all__ = [
"Array2D",
"Array3D",
"Array4D",
"ArrayND",
"MArray",
"MArray0D",
"MArray1D",
"MArray2D",
"MArray3D",
"MArrayND",
"Matrix",
"StringArray",
"StringArray0D",
Expand All @@ -33,7 +31,8 @@ __all__ = [

###

_RankT = TypeVar("_RankT", bound=Shape, default=Shape)
# TODO: use `Shape` instead of `AnyShape` once python/mypy#19110 is fixed
_RankT = TypeVar("_RankT", bound=AnyShape, default=Shape)
_ScalarT = TypeVar("_ScalarT", bound=np.generic, default=Any)
_NaT = TypeVar("_NaT", default=Never)

Expand All @@ -45,7 +44,6 @@ Array1D = TypeAliasType("Array1D", np.ndarray[Rank1, np.dtype[_ScalarT]], type_p
Array2D = TypeAliasType("Array2D", np.ndarray[Rank2, np.dtype[_ScalarT]], type_params=(_ScalarT,))
Array3D = TypeAliasType("Array3D", np.ndarray[Rank3, np.dtype[_ScalarT]], type_params=(_ScalarT,))
Array4D = TypeAliasType("Array4D", np.ndarray[Rank4, np.dtype[_ScalarT]], type_params=(_ScalarT,))
ArrayND = TypeAliasType("ArrayND", np.ndarray[Rank, np.dtype[_ScalarT]], type_params=(_ScalarT,))

###

Expand All @@ -58,7 +56,6 @@ MArray0D = TypeAliasType("MArray0D", np.ma.MaskedArray[Rank0, np.dtype[_ScalarT]
MArray1D = TypeAliasType("MArray1D", np.ma.MaskedArray[Rank1, np.dtype[_ScalarT]], type_params=(_ScalarT,))
MArray2D = TypeAliasType("MArray2D", np.ma.MaskedArray[Rank2, np.dtype[_ScalarT]], type_params=(_ScalarT,))
MArray3D = TypeAliasType("MArray3D", np.ma.MaskedArray[Rank3, np.dtype[_ScalarT]], type_params=(_ScalarT,))
MArrayND = TypeAliasType("MArrayND", np.ma.MaskedArray[Rank, np.dtype[_ScalarT]], type_params=(_ScalarT, _RankT))

###

Expand Down Expand Up @@ -89,6 +86,6 @@ StringArray3D = TypeAliasType(
)
StringArrayND = TypeAliasType(
"StringArrayND",
np.ndarray[Rank, np.dtypes.StringDType[_NaT]],
np.ndarray[Shape, np.dtypes.StringDType[_NaT]],
type_params=(_NaT,),
)
2 changes: 1 addition & 1 deletion src/_numtype/_nep50.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ _ScalarOutT = TypeVar("_ScalarOutT", bound=_ScalarOut, default=Any)
_ScalarOutT_co = TypeVar("_ScalarOutT_co", bound=_ScalarOut, covariant=True)
_ScalarOutT_contra = TypeVar("_ScalarOutT_contra", bound=_ScalarOut, contravariant=True)

_ShapeT = TypeVar("_ShapeT", bound=_shape.Shape, default=_shape.Shape)
_ShapeT = TypeVar("_ShapeT", bound=_shape.Shape, default=Any)
_ShapeT_co = TypeVar("_ShapeT_co", bound=_shape.Shape, covariant=True)

###
Expand Down
115 changes: 73 additions & 42 deletions src/_numtype/_rank.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
from typing import Any, Generic, Protocol, Self, TypeAlias, final, type_check_only
from typing_extensions import TypeAliasType, TypeVar
from typing_extensions import TypeAliasType, TypeVar, TypeVarTuple, override

from ._shape import Shape, Shape0, Shape0N, Shape1, Shape1N, Shape2, Shape2N, Shape3, Shape3N, Shape4, Shape4N
from ._shape import AnyShape, Shape, Shape0, Shape1, Shape1N, Shape2, Shape2N, Shape3, Shape3N, Shape4, Shape4N

__all__ = [
"HasInnerShape",
"HasRankGE",
"HasRankLE",
"Rank",
Expand All @@ -21,56 +22,71 @@ __all__ = [

###

_Shape00: TypeAlias = Shape0
_Shape01: TypeAlias = _Shape00 | Shape1
_Shape01: TypeAlias = Shape0 | Shape1
_Shape02: TypeAlias = _Shape01 | Shape2
_Shape03: TypeAlias = _Shape02 | Shape3
_Shape04: TypeAlias = _Shape03 | Shape4

###

_UpperT = TypeVar("_UpperT", bound=Shape)
_LowerT = TypeVar("_LowerT", bound=Shape)
# TODO(jorenham): remove `| Rank0 | Rank` once python/mypy#19110 is fixed
_UpperT = TypeVar("_UpperT", bound=Shape | Rank0 | Rank)
_LowerT = TypeVar("_LowerT", bound=Shape | Rank0 | Rank)
_RankT = TypeVar("_RankT", bound=Shape, default=Any)

# TODO(jorenham): remove `| Rank0 | Rank` once python/mypy#19110 is fixed
_RankLE: TypeAlias = _CanBroadcast[Any, _UpperT, _RankT] | Shape0 | Rank0 | Rank
# TODO(jorenham): remove `| Rank` once python/mypy#19110 is fixed
_RankGE: TypeAlias = _CanBroadcast[_LowerT, Any, _RankT] | _LowerT | Rank

HasRankLE = TypeAliasType(
"HasRankLE",
_HasShape[Shape0 | _HasOwnShape[_UpperT] | _CanBroadcast[Any, _UpperT, _RankT]],
_HasInnerShape[_RankLE[_UpperT, _RankT]],
type_params=(_UpperT, _RankT),
)
HasRankGE = TypeAliasType(
"HasRankGE",
_HasShape[_LowerT | _CanBroadcast[_LowerT, Any, _RankT]],
_HasInnerShape[_RankGE[_LowerT, _RankT]],
type_params=(_LowerT, _RankT),
)

###
_ShapeT = TypeVar("_ShapeT", bound=Shape)

_ShapeT_co = TypeVar("_ShapeT_co", bound=Shape | _HasOwnShape | _CanBroadcast, covariant=True)
# for unwrapping potential rank types as shape tuples
HasInnerShape = TypeAliasType(
"HasInnerShape",
_HasInnerShape[_HasOwnShape[Any, _ShapeT]],
type_params=(_ShapeT,),
)

@type_check_only
class _HasShape(Protocol[_ShapeT_co]):
@property
def shape(self, /) -> _ShapeT_co: ...
###

_ShapeLikeT_co = TypeVar("_ShapeLikeT_co", bound=Shape | _HasOwnShape | _CanBroadcast[Any, Any], covariant=True)

_FromT_contra = TypeVar("_FromT_contra", default=Any, contravariant=True)
_ToT_contra = TypeVar("_ToT_contra", bound=Shape, default=Any, contravariant=True)
_FromT_contra = TypeVar("_FromT_contra", contravariant=True)
_ToT_contra = TypeVar("_ToT_contra", bound=tuple[Any, ...], contravariant=True)
_EquivT_co = TypeVar("_EquivT_co", bound=Shape, default=Any, covariant=True)

# __broadcast__ is the type-check-only interface order of ranks
@final
@type_check_only
class _CanBroadcast(Protocol[_FromT_contra, _ToT_contra, _EquivT_co]):
def __broadcast__(self, from_: _FromT_contra, to: _ToT_contra, /) -> _EquivT_co: ...

# __inner_shape__ is similar to `shape`, but directly exposes the `Rank` type.
@final
@type_check_only
class _HasInnerShape(Protocol[_ShapeLikeT_co]):
@property
def __inner_shape__(self, /) -> _ShapeLikeT_co: ...

_OwnShapeT_contra = TypeVar("_OwnShapeT_contra", bound=tuple[Any, ...], default=Any, contravariant=True)
_OwnShapeT_co = TypeVar("_OwnShapeT_co", bound=Shape, default=_OwnShapeT_contra, covariant=True)

# This double shape-type parameter is a sneaky way to annotate a doubly-bound nominal type range,
# e.g. `_HasOwnShape[Shape2N, Shape0N]` accepts `Shape2N`, `Shape1N`, and `Shape0N`, but
# rejects `Shape3N` and `Shape1`. Besides brevity, it also works around several mypy bugs that
# are related to "unions vs joins".

_OwnShapeT_contra = TypeVar("_OwnShapeT_contra", bound=Shape, default=Any, contravariant=True)
_OwnShapeT_co = TypeVar("_OwnShapeT_co", bound=Shape, default=_OwnShapeT_contra, covariant=True)
_OwnShapeT = TypeVar("_OwnShapeT", bound=tuple[Any, ...], default=Any)

@final
@type_check_only
class _HasOwnShape(Protocol[_OwnShapeT_contra, _OwnShapeT_co]):
Expand All @@ -79,59 +95,74 @@ class _HasOwnShape(Protocol[_OwnShapeT_contra, _OwnShapeT_co]):
###
# TODO(jorenham): embed the array-like types, e.g. `Sequence[Sequence[T]]`

@type_check_only
class _BaseRank(Generic[_FromT_contra, _OwnShapeT, _ToT_contra]):
def __broadcast__(self, from_: _FromT_contra, to: _ToT_contra, /) -> Self: ...
def __own_shape__(self, shape: _OwnShapeT, /) -> _OwnShapeT: ...
_Ts = TypeVarTuple("_Ts") # should only contain `int`s

# https://github.com/python/mypy/issues/19093
@type_check_only
class _BaseRankM(
_BaseRank[_FromT_contra | _HasOwnShape[_ToT_contra, Shape], _OwnShapeT, _ToT_contra],
Generic[_FromT_contra, _OwnShapeT, _ToT_contra],
): ...
class BaseRank(tuple[*_Ts], Generic[*_Ts]):
def __broadcast__(self, from_: tuple[*_Ts], to: tuple[*_Ts], /) -> Self: ...
def __own_shape__(self, shape: tuple[*_Ts], /) -> tuple[*_Ts]: ...

@final
@type_check_only
class Rank0(_BaseRankM[_Shape00, Shape0, Shape0N], tuple[()]): ...
class Rank0(BaseRank[()]):
@override
def __broadcast__(self, from_: Shape0 | _HasOwnShape[Shape, Any], to: Shape, /) -> Self: ...

@final
@type_check_only
class Rank1(_BaseRankM[_Shape01, Shape1, Shape1N], tuple[int]): ...
class Rank1(BaseRank[int]):
@override
def __broadcast__(self, from_: _Shape01 | _HasOwnShape[Shape1N, Any], to: Shape1N, /) -> Self: ...

@final
@type_check_only
class Rank2(_BaseRankM[_Shape02, Shape2, Shape2N], tuple[int, int]): ...
class Rank2(BaseRank[int, int]):
@override
def __broadcast__(self, from_: _Shape02 | _HasOwnShape[Shape2N, Any], to: Shape2N, /) -> Self: ...

@final
@type_check_only
class Rank3(_BaseRankM[_Shape03, Shape3, Shape3N], tuple[int, int, int]): ...
class Rank3(BaseRank[int, int, int]):
@override
def __broadcast__(self, from_: _Shape03 | _HasOwnShape[Shape3N, Any], to: Shape3N, /) -> Self: ...

@final
@type_check_only
class Rank4(_BaseRankM[_Shape04, Shape4, Shape4N], tuple[int, int, int, int]): ...
class Rank4(BaseRank[int, int, int, int]):
@override
def __broadcast__(self, from_: _Shape04 | _HasOwnShape[Shape4N, Any], to: Shape4N, /) -> Self: ...

# this emulates `AnyOf`, rather than a `Union`.
@type_check_only
class _BaseRankMToN(_BaseRank[Shape0N, _OwnShapeT, _OwnShapeT], Generic[_OwnShapeT]): ...
# these emulates `AnyOf` (gradual union), rather than a `Union`.

@final
@type_check_only
class Rank(_BaseRankMToN[Shape0N], tuple[int, ...]): ...
class Rank(BaseRank[*tuple[int, ...]]):
@override
def __broadcast__(self, from_: AnyShape, to: tuple[*_Ts], /) -> Self: ...

@final
@type_check_only
class Rank1N(_BaseRankMToN[Shape1N], tuple[int, *tuple[int, ...]]): ...
class Rank1N(BaseRank[int, *tuple[int, ...]]):
@override
def __broadcast__(self, from_: AnyShape, to: Shape1N, /) -> Self: ...

@final
@type_check_only
class Rank2N(_BaseRankMToN[Shape2N], tuple[int, int, *tuple[int, ...]]): ...
class Rank2N(BaseRank[int, int, *tuple[int, ...]]):
@override
def __broadcast__(self, from_: AnyShape, to: Shape2N, /) -> Self: ...

@final
@type_check_only
class Rank3N(_BaseRankMToN[Shape3N], tuple[int, int, int, *tuple[int, ...]]): ...
class Rank3N(BaseRank[int, int, int, *tuple[int, ...]]):
@override
def __broadcast__(self, from_: AnyShape, to: Shape3N, /) -> Self: ...

@final
@type_check_only
class Rank4N(_BaseRankMToN[Shape4N], tuple[int, int, int, int, *tuple[int, ...]]): ...
class Rank4N(BaseRank[int, int, int, int, *tuple[int, ...]]):
@override
def __broadcast__(self, from_: AnyShape, to: Shape4N, /) -> Self: ...

Rank0N: TypeAlias = Rank
5 changes: 4 additions & 1 deletion src/_numtype/_shape.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
from typing import TypeAlias
from typing import Any, TypeAlias
from typing_extensions import TypeAliasType

__all__ = [
"AnyShape",
"Shape",
"Shape0",
"Shape0N",
Expand All @@ -16,8 +17,10 @@ __all__ = [
"ShapeN",
]

AnyShape = TypeAliasType("AnyShape", tuple[Any, ...])
Shape = TypeAliasType("Shape", tuple[int, ...])

# TODO: remove `| Rank0` once python/mypy#19110 is fixed
Shape0 = TypeAliasType("Shape0", tuple[()])
Shape1 = TypeAliasType("Shape1", tuple[int])
Shape2 = TypeAliasType("Shape2", tuple[int, int])
Expand Down
Loading
0