Skip to content

Commit 2f62ddc

Browse files
authored
Merge pull request #198 from jorenham/numpy-3d-aliases
add 3-d aliases for arrays and array-likes in `optype.numpy`
2 parents 54c72fa + 19717b6 commit 2f62ddc

File tree

3 files changed

+86
-38
lines changed

3 files changed

+86
-38
lines changed

README.md

Lines changed: 10 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2529,29 +2529,27 @@ type AtMost3D[
25292529
Similar to the `numpy._typing._ArrayLike{}_co` *coercible array-like* types,
25302530
`optype.numpy` provides the `optype.numpy.To{}ND`. Unlike the ones in `numpy`, these
25312531
don't accept "bare" scalar types (the `__len__` method is required).
2532-
Additionally, there are the `To{}1D` and `To{}2D` for vector-likes and matrix-likes, and
2533-
the `To{}` aliases for "bare" scalar types.
2532+
Additionally, there are the `To{}1D`, `To{}2D`, and `To{}3D` for vector-likes,
2533+
matrix-likes, and cuboid-likes, and the `To{}` aliases for "bare" scalar types.
25342534
25352535
<table>
25362536
<tr>
25372537
<th align="center" colspan="2">scalar types</th>
25382538
<th align="center">scalar-like</th>
2539-
<th align="center">1-d array-like</th>
2540-
<th align="center">2-d array-like</th>
2541-
<th align="center">n-d array-like</th>
2539+
<th align="center"><code>{1,2,3}</code>-d array-like</th>
2540+
<th align="center"><code>N</code>-d array-like</th>
25422541
</tr>
25432542
<tr>
25442543
<th align="center"><code>builtins</code></th>
25452544
<th align="center"><code>numpy</code></th>
25462545
<th align="center" colspan="4"><code>optype.numpy</code></th>
25472546
</tr>
2548-
<tr><td colspan="6"></td></tr>
2547+
<tr><td colspan="5"></td></tr>
25492548
<tr>
25502549
<td align="left"><code>bool</code></td>
25512550
<td align="left"><code>bool_</code></td>
25522551
<td align="left"><code>ToBool</code></td>
2553-
<td align="left"><code>ToBool1D</code></td>
2554-
<td align="left"><code>ToBool2D</code></td>
2552+
<td align="left"><code>ToBool{1,2,3}D</code></td>
25552553
<td align="left"><code>ToBoolND</code></td>
25562554
</tr>
25572555
<tr><td colspan="6"></td></tr>
@@ -2562,8 +2560,7 @@ the `To{}` aliases for "bare" scalar types.
25622560
<code>| bool_</code>
25632561
</td>
25642562
<td align="left"><code>ToInt</code></td>
2565-
<td align="left"><code>ToInt1D</code></td>
2566-
<td align="left"><code>ToInt2D</code></td>
2563+
<td align="left"><code>ToInt{1,2,3}D</code></td>
25672564
<td align="left"><code>ToIntND</code></td>
25682565
</tr>
25692566
<tr><td colspan="6"></td></tr>
@@ -2578,8 +2575,7 @@ the `To{}` aliases for "bare" scalar types.
25782575
<code>| bool_</code>
25792576
</td>
25802577
<td align="left"><code>ToFloat</code></td>
2581-
<td align="left"><code>ToFloat1D</code></td>
2582-
<td align="left"><code>ToFloat2D</code></td>
2578+
<td align="left"><code>ToFloat{1,2,3}D</code></td>
25832579
<td align="left"><code>ToFloatND</code></td>
25842580
</tr>
25852581
<tr><td colspan="6"></td></tr>
@@ -2594,8 +2590,7 @@ the `To{}` aliases for "bare" scalar types.
25942590
<code>| bool_</code>
25952591
</td>
25962592
<td align="left"><code>ToComplex</code></td>
2597-
<td align="left"><code>ToComplex1D</code></td>
2598-
<td align="left"><code>ToComplex2D</code></td>
2593+
<td align="left"><code>ToComplex{1,2,3}D</code></td>
25992594
<td align="left"><code>ToComplexND</code></td>
26002595
</tr>
26012596
<tr><td colspan="6"></td></tr>
@@ -2609,8 +2604,7 @@ the `To{}` aliases for "bare" scalar types.
26092604
</td>
26102605
<td align="left"><code>generic</code></td>
26112606
<td align="left"><code>ToScalar</code></td>
2612-
<td align="left"><code>ToArray1D</code></td>
2613-
<td align="left"><code>ToArray2D</code></td>
2607+
<td align="left"><code>ToArray{1,2,3}D</code></td>
26142608
<td align="left"><code>ToArrayND</code></td>
26152609
</tr>
26162610
</table>

optype/numpy/_array.py

Lines changed: 57 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,6 @@
99
import optype.numpy._compat as _x
1010
from optype._core._utils import set_module
1111

12-
from ._dtype import DType
13-
from ._shape import AtLeast0D
14-
1512

1613
if sys.version_info >= (3, 13):
1714
from typing import Protocol, Self, TypeAliasType, TypeVar, runtime_checkable
@@ -33,8 +30,12 @@
3330
"Array1D",
3431
"Array1D",
3532
"Array2D",
33+
"Array3D",
3634
"ArrayND",
3735
"CanArray",
36+
"CanArray1D",
37+
"CanArray2D",
38+
"CanArray3D",
3839
"CanArrayFinalize",
3940
"CanArrayND",
4041
"CanArrayWrap",
@@ -43,10 +44,20 @@
4344
]
4445

4546

46-
_NDT = TypeVar("_NDT", bound=AtLeast0D, default=AtLeast0D)
47-
_NDT_co = TypeVar("_NDT_co", bound=AtLeast0D, default=AtLeast0D, covariant=True)
48-
_DTT = TypeVar("_DTT", bound=DType, default=DType)
49-
_DTT_co = TypeVar("_DTT_co", bound=DType, default=DType, covariant=True)
47+
_NDT = TypeVar("_NDT", bound=tuple[int, ...], default=tuple[int, ...])
48+
_NDT_co = TypeVar(
49+
"_NDT_co",
50+
bound=tuple[int, ...],
51+
default=tuple[int, ...],
52+
covariant=True,
53+
)
54+
_DTT = TypeVar("_DTT", bound=np.dtype[np.generic], default=np.dtype[np.generic])
55+
_DTT_co = TypeVar(
56+
"_DTT_co",
57+
bound=np.dtype[np.generic],
58+
default=np.dtype[np.generic],
59+
covariant=True,
60+
)
5061
_SCT = TypeVar("_SCT", bound=np.generic, default=np.generic)
5162
_SCT_co = TypeVar("_SCT_co", bound=np.generic, default=np.generic, covariant=True)
5263

@@ -115,6 +126,45 @@ class CanArray(Protocol[_NDT, _DTT_co]):
115126
def __array__(self, /) -> np.ndarray[_NDT, _DTT_co]: ...
116127

117128

129+
@runtime_checkable
130+
@set_module("optype.numpy")
131+
class CanArrayND(Protocol[_SCT_co]):
132+
"""
133+
Similar to `optype.numpy.CanArray`, but must be sized (i.e. excludes scalars),
134+
and is parameterized by only the scalar type (instead of the shape and dtype).
135+
"""
136+
137+
def __len__(self, /) -> int: ...
138+
def __array__(self, /) -> np.ndarray[tuple[int, ...], np.dtype[_SCT_co]]: ...
139+
140+
141+
@runtime_checkable
142+
@set_module("optype.numpy")
143+
class CanArray1D(Protocol[_SCT_co]):
144+
"""The 1-d variant of `optype.numpy.CanArrayND`."""
145+
146+
def __len__(self, /) -> int: ...
147+
def __array__(self, /) -> np.ndarray[tuple[int], np.dtype[_SCT_co]]: ...
148+
149+
150+
@runtime_checkable
151+
@set_module("optype.numpy")
152+
class CanArray2D(Protocol[_SCT_co]):
153+
"""The 2-d variant of `optype.numpy.CanArrayND`."""
154+
155+
def __len__(self, /) -> int: ...
156+
def __array__(self, /) -> np.ndarray[tuple[int, int], np.dtype[_SCT_co]]: ...
157+
158+
159+
@runtime_checkable
160+
@set_module("optype.numpy")
161+
class CanArray3D(Protocol[_SCT_co]):
162+
"""The 2-d variant of `optype.numpy.CanArrayND`."""
163+
164+
def __len__(self, /) -> int: ...
165+
def __array__(self, /) -> np.ndarray[tuple[int, int, int], np.dtype[_SCT_co]]: ...
166+
167+
118168
# this is almost always a `ndarray`, but setting a `bound` might break in some
119169
# edge cases
120170
_T_contra = TypeVar("_T_contra", contravariant=True, default=object)
@@ -169,15 +219,3 @@ def __array_interface__(self, /) -> _ArrayInterfaceT_co: ...
169219
class HasArrayPriority(Protocol):
170220
@property
171221
def __array_priority__(self, /) -> float: ...
172-
173-
174-
@runtime_checkable
175-
@set_module("optype.numpy")
176-
class CanArrayND(Protocol[_SCT_co]):
177-
"""
178-
Similar to `optype.numpy.CanArray`, but must be sized (i.e. excludes scalars),
179-
and is parameterized by only the scalar type (instead of the shape and dtype).
180-
"""
181-
182-
def __len__(self, /) -> int: ...
183-
def __array__(self, /) -> np.ndarray[tuple[int, ...], np.dtype[_SCT_co]]: ...

optype/numpy/_to.py

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from __future__ import annotations
22

33
import sys
4-
from collections.abc import Sequence
4+
from collections.abc import Sequence as Seq
55
from typing import TypeAlias, TypeVar
66

77
import numpy as np
@@ -20,22 +20,27 @@
2020
__all__ = [
2121
"ToArray1D",
2222
"ToArray2D",
23+
"ToArray3D",
2324
"ToArrayND",
2425
"ToBool",
2526
"ToBool1D",
2627
"ToBool2D",
28+
"ToBool3D",
2729
"ToBoolND",
2830
"ToComplex",
2931
"ToComplex1D",
3032
"ToComplex2D",
33+
"ToComplex3D",
3134
"ToComplexND",
3235
"ToFloat",
3336
"ToFloat1D",
3437
"ToFloat2D",
38+
"ToFloat3D",
3539
"ToFloatND",
3640
"ToInt",
3741
"ToInt1D",
3842
"ToInt2D",
43+
"ToInt3D",
3944
"ToIntND",
4045
"ToScalar",
4146
]
@@ -44,36 +49,47 @@
4449
ST = TypeVar("ST", bound=np.generic)
4550

4651

47-
_To1D: TypeAlias = CanArrayND[ST] | Sequence[T | ST]
48-
_To2D: TypeAlias = CanArrayND[ST] | Sequence[_To1D[ST, T]]
52+
_To1D: TypeAlias = CanArrayND[ST] | Seq[ST | T]
53+
_To2D: TypeAlias = CanArrayND[ST] | Seq[CanArrayND[ST]] | Seq[Seq[ST | T]]
54+
_To3D: TypeAlias = (
55+
CanArrayND[ST]
56+
| Seq[CanArrayND[ST]]
57+
| Seq[Seq[CanArrayND[ST]]]
58+
| Seq[Seq[Seq[ST | T]]]
59+
)
4960
# recursive sequence type cannot be used due to a mypy bug:
5061
# https://github.com/python/mypy/issues/18184
5162
_ToND: TypeAlias = CanArrayND[ST] | SequenceND[T | ST] | SequenceND[CanArrayND[ST]]
5263

5364
ToScalar: TypeAlias = complex | bytes | str | np.generic
5465
ToArray1D: TypeAlias = _To1D[np.generic, complex | bytes | str]
5566
ToArray2D: TypeAlias = _To2D[np.generic, complex | bytes | str]
67+
ToArray3D: TypeAlias = _To3D[np.generic, complex | bytes | str]
5668
ToArrayND: TypeAlias = _ToND[np.generic, complex | bytes | str]
5769

5870
ToBool: TypeAlias = bool | np.bool_
5971
ToBool1D: TypeAlias = _To1D[np.bool_, bool]
6072
ToBool2D: TypeAlias = _To2D[np.bool_, bool]
73+
ToBool3D: TypeAlias = _To3D[np.bool_, bool]
6174
ToBoolND: TypeAlias = _ToND[np.bool_, bool]
6275

6376
integer_co = TypeAliasType("integer_co", integer | np.bool_) # type: ignore[no-any-explicit]
6477
ToInt: TypeAlias = int | integer_co
6578
ToInt1D: TypeAlias = _To1D[integer_co, int]
6679
ToInt2D: TypeAlias = _To2D[integer_co, int]
80+
ToInt3D: TypeAlias = _To3D[integer_co, int]
6781
ToIntND: TypeAlias = _ToND[integer_co, int]
6882

6983
floating_co = TypeAliasType("floating_co", floating | integer | np.bool_) # type: ignore[no-any-explicit]
7084
ToFloat: TypeAlias = float | floating_co
7185
ToFloat1D: TypeAlias = _To1D[floating_co, float]
7286
ToFloat2D: TypeAlias = _To2D[floating_co, float]
87+
ToFloat3D: TypeAlias = _To3D[floating_co, float]
7388
ToFloatND: TypeAlias = _ToND[floating_co, float]
7489

7590
complexfloating_co = TypeAliasType("complexfloating_co", number | np.bool_) # type: ignore[no-any-explicit]
7691
ToComplex: TypeAlias = complex | complexfloating_co
7792
ToComplex1D: TypeAlias = _To1D[complexfloating_co, complex]
7893
ToComplex2D: TypeAlias = _To2D[complexfloating_co, complex]
94+
ToComplex3D: TypeAlias = _To3D[complexfloating_co, complex]
7995
ToComplexND: TypeAlias = _ToND[complexfloating_co, complex]

0 commit comments

Comments
 (0)