🎨 prefer _numtype.Shape* over int tuples as shape type (#532)

Author: jorenham

src/_numtype/__init__.pyi

@@ -684,10 +684,18 @@ ToObject_2nd = TypeAliasType("ToObject_2nd", _ToArray2_2nd[np.object_, _PyObject
 ToObject_3nd = TypeAliasType("ToObject_3nd", _ToArray2_3nd[np.object_, _PyObject])
 
 # StringDType
-ToString_nd = TypeAliasType("ToString_nd", _CanStringArray[Shape, _NaT0], type_params=(_NaT0,))
+ToString_nd = TypeAliasType("ToString_nd", SequenceND[_CanStringArray[Shape, _NaT0]], type_params=(_NaT0,))
 ToString_1ds = TypeAliasType("ToString_1ds", _CanStringArray[Shape1, _NaT0], type_params=(_NaT0,))
-ToString_2ds = TypeAliasType("ToString_2ds", _CanStringArray[Shape2, _NaT0], type_params=(_NaT0,))
-ToString_3ds = TypeAliasType("ToString_3ds", _CanStringArray[Shape3, _NaT0], type_params=(_NaT0,))
+ToString_2ds = TypeAliasType(
+    "ToString_2ds",
+    _CanStringArray[Shape2, _NaT0] | Sequence[ToString_1ds[_NaT0]],
+    type_params=(_NaT0,),
+)
+ToString_3ds = TypeAliasType(
+    "ToString_3ds",
+    _CanStringArray[Shape3, _NaT0] | Sequence[ToString_2ds[_NaT0]],
+    type_params=(_NaT0,),
+)
 ToString_1nd = TypeAliasType("ToString_1nd", _CanStringArray[Shape1_, _NaT0], type_params=(_NaT0,))
 ToString_2nd = TypeAliasType("ToString_2nd", _CanStringArray[Shape2_, _NaT0], type_params=(_NaT0,))
 ToString_3nd = TypeAliasType("ToString_3nd", _CanStringArray[Shape3_, _NaT0], type_params=(_NaT0,))

src/numpy-stubs/@test/runtime/legacy/arrayterator.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 import numpy as np
+import numpy.typing as npt
 
-AR_i8: np.ndarray[tuple[int, ...], np.dtype[np.int_]] = np.arange(10)
+AR_i8: npt.NDArray[np.int_] = np.arange(10)
 ar_iter = np.lib.Arrayterator(AR_i8)
 
 ar_iter.var

src/numpy-stubs/@test/runtime/legacy/lib_user_array.py

@@ -2,21 +2,26 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 from numpy.lib.user_array import container
 
+if TYPE_CHECKING:
+    import _numtype as _nt
+
 N = 10_000
 W = H = int(N**0.5)
 
-a: np.ndarray[tuple[int, int], np.dtype[np.int32]]
-ua: container[tuple[int, int], np.dtype[np.int32]]
+a: np.ndarray[_nt.Shape2, np.dtype[np.int32]]
+ua: container[_nt.Shape2, np.dtype[np.int32]]
 
 a = np.arange(N, dtype=np.int32).reshape(W, H)
 ua = container(a)
 
-ua_small: container[tuple[int, int], np.dtype[np.int32]] = ua[:3, :5]
+ua_small: container[_nt.Shape2, np.dtype[np.int32]] = ua[:3, :5]
 ua_small[0, 0] = 10
 
-ua_bool: container[tuple[int, int], np.dtype[np.bool]] = ua_small > 1
+ua_bool: container[_nt.Shape2, np.dtype[np.bool]] = ua_small > 1
 
 # shape: tuple[int, int] = np.shape(ua)

src/numpy-stubs/@test/runtime/legacy/ndarray_shape_manipulation.py

@@ -1,8 +1,13 @@
-from typing import cast
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 
-nd1 = cast("np.ndarray[tuple[int, int], np.dtype[np.intp]]", np.array([[1, 2], [3, 4]]))
+if TYPE_CHECKING:
+    import _numtype as _nt
+
+nd1 = cast("np.ndarray[_nt.Shape2, np.dtype[np.intp]]", np.array([[1, 2], [3, 4]]))
 
 # reshape
 nd1.reshape(4)

src/numpy-stubs/@test/runtime/legacy/recfunctions.py

@@ -1,12 +1,15 @@
 """These tests are based on the doctests from `numpy/lib/recfunctions.py`."""
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 from typing_extensions import assert_type
 
 import numpy as np
 import numpy.typing as npt
 from numpy.lib import recfunctions as rfn
 
+if TYPE_CHECKING:
+    import _numtype as _nt  # noqa: F401, PLC2701
+
 
 def test_recursive_fill_fields() -> None:
     a: npt.NDArray[np.void] = np.array(
@@ -15,7 +18,7 @@ def test_recursive_fill_fields() -> None:
     )
     b = np.zeros((int(3),), dtype=a.dtype)
     out = rfn.recursive_fill_fields(a, b)
-    assert_type(out, np.ndarray[tuple[int], np.dtype[np.void]])
+    assert_type(out, "np.ndarray[_nt.Shape1, np.dtype[np.void]]")
 
 
 def test_get_names() -> None:
@@ -57,7 +60,7 @@ def test_merge_arrays() -> None:
             np.ones((int(2),), np.int_),
             np.ones((int(3),), np.float64),
         )),
-        np.recarray[tuple[int], np.dtype[np.void]],
+        "np.recarray[_nt.Shape1, np.dtype[np.void]]",
     )
 
 
@@ -67,15 +70,15 @@ def test_drop_fields() -> None:
 
     assert_type(
         rfn.drop_fields(a, "a"),
-        np.ndarray[tuple[int], np.dtype[np.void]],
+        "np.ndarray[_nt.Shape1, np.dtype[np.void]]",
     )
     assert_type(
         rfn.drop_fields(a, "a", asrecarray=True),
-        np.rec.recarray[tuple[int], np.dtype[np.void]],
+        "np.rec.recarray[_nt.Shape1, np.dtype[np.void]]",
     )
     assert_type(
         rfn.rec_drop_fields(a, "a"),
-        np.rec.recarray[tuple[int], np.dtype[np.void]],
+        "np.rec.recarray[_nt.Shape1, np.dtype[np.void]]",
     )
 
 
@@ -85,7 +88,7 @@ def test_rename_fields() -> None:
 
     assert_type(
         rfn.rename_fields(a, {"a": "A", "b_b": "B_B"}),
-        np.ndarray[tuple[int], np.dtype[np.void]],
+        "np.ndarray[_nt.Shape1, np.dtype[np.void]]",
     )
 
 
@@ -96,7 +99,7 @@ def test_repack_fields() -> None:
     assert_type(rfn.repack_fields(dt.type(0)), np.void)
     assert_type(
         rfn.repack_fields(np.ones((3,), dtype=dt)),
-        np.ndarray[tuple[int], np.dtype[np.void]],
+        "np.ndarray[_nt.Shape1, np.dtype[np.void]]",
     )
 
 
@@ -115,38 +118,38 @@ def test_apply_along_fields() -> None:
     b = np.ones(4, dtype=[("x", "i4"), ("y", "f4"), ("z", "f8")])
     assert_type(
         rfn.apply_along_fields(np.mean, b),
-        np.ndarray[tuple[int], np.dtype[np.void]],
+        "np.ndarray[_nt.Shape1, np.dtype[np.void]]",
     )
 
 
 def test_assign_fields_by_name() -> None:
     b = np.ones(4, dtype=[("x", "i4"), ("y", "f4"), ("z", "f8")])
     assert_type(
         rfn.apply_along_fields(np.mean, b),
-        np.ndarray[tuple[int], np.dtype[np.void]],
+        "np.ndarray[_nt.Shape1, np.dtype[np.void]]",
     )
 
 
 def test_require_fields() -> None:
     a = np.ones(4, dtype=[("a", "i4"), ("b", "f8"), ("c", "u1")])
     assert_type(
         rfn.require_fields(a, [("b", "f4"), ("c", "u1")]),
-        np.ndarray[tuple[int], np.dtype[np.void]],
+        "np.ndarray[_nt.Shape1, np.dtype[np.void]]",
     )
 
 
 def test_stack_arrays() -> None:
     x = np.zeros((int(2),), np.int32)
     assert_type(
         rfn.stack_arrays(x),
-        np.ndarray[tuple[int], np.dtype[np.int32]],
+        "np.ndarray[_nt.Shape1, np.dtype[np.int32]]",
     )
 
     z = np.ones((int(2),), [("A", "|S3"), ("B", float)])
     zz = np.ones((int(2),), [("A", "|S3"), ("B", np.float64), ("C", np.float64)])
     assert_type(
         rfn.stack_arrays((z, zz)),
-        np.ma.MaskedArray[tuple[int, ...], np.dtype[np.void]],
+        "np.ma.MaskedArray[_nt.Shape, np.dtype[np.void]]",
     )
 
 

src/numpy-stubs/@test/runtime/legacy/shape.py

@@ -1,7 +1,10 @@
-from typing import Any, NamedTuple
+from typing import TYPE_CHECKING, Any, NamedTuple
 
 import numpy as np
 
+if TYPE_CHECKING:
+    import _numtype as _nt
+
 
 # Subtype of tuple[int, int]
 class XYGrid(NamedTuple):
@@ -10,7 +13,7 @@ class XYGrid(NamedTuple):
 
 
 # Test variance of _ShapeType_co
-def accepts_2d(a: np.ndarray[tuple[int, int], Any]) -> None:
+def accepts_2d(a: "np.ndarray[_nt.Shape2, Any]") -> None:
     return None