Merge pull request #548 from numpy/prefer-_nt.Array

Author: jorenham

src/numpy-stubs/@test/static/accept/chararray.pyi

@@ -1,10 +1,15 @@
-from typing import TypeAlias, assert_type
+from typing import Any, TypeAlias, assert_type
+from typing_extensions import TypeVar
 
 import _numtype as _nt
 import numpy as np
 
-_BytesArray: TypeAlias = np.char.chararray[_nt.Shape, np.dtype[np.bytes_]]
-_StrArray: TypeAlias = np.char.chararray[_nt.Shape, np.dtype[np.str_]]
+_ShapeT = TypeVar("_ShapeT", bound=_nt.Shape, default=_nt.Shape)
+
+_BytesArray: TypeAlias = np.char.chararray[_ShapeT, np.dtype[np.bytes_]]
+_StrArray: TypeAlias = np.char.chararray[_ShapeT, np.dtype[np.str_]]
+
+###
 
 AR_U: _StrArray
 AR_S: _BytesArray
@@ -85,8 +90,8 @@ assert_type(AR_S.title(), _BytesArray)
 assert_type(AR_U.upper(), _StrArray)
 assert_type(AR_S.upper(), _BytesArray)
 
-assert_type(AR_U.zfill(5), _StrArray)
-assert_type(AR_S.zfill([2, 3, 4]), _BytesArray)
+assert_type(AR_U.zfill(5), _StrArray[Any])
+assert_type(AR_S.zfill([2, 3, 4]), _BytesArray[Any])
 
 assert_type(AR_U.count("a", start=[1, 2, 3]), _nt.Array[np.int_])
 assert_type(AR_S.count([b"a", b"b", b"c"], end=9), _nt.Array[np.int_])

src/numpy-stubs/__init__.pyi

@@ -1,11 +1,12 @@
 from collections.abc import Callable
 from contextlib import _GeneratorContextManager
-from typing import Any, Final, Literal, SupportsIndex, TypeAlias, TypedDict, overload, type_check_only
+from typing import Final, Literal, SupportsIndex, TypeAlias, TypedDict, overload, type_check_only
 from typing_extensions import deprecated
 
+import _numtype as _nt
 import numpy as np
 from numpy._globals import _NoValueType
-from numpy._typing import NDArray, _CharLike_co, _FloatLike_co
+from numpy._typing import _CharLike_co, _FloatLike_co
 
 __all__ = [
     "array2string",
@@ -25,7 +26,7 @@ _Sign: TypeAlias = Literal["-", "+", " "]
 _Trim: TypeAlias = Literal["k", ".", "0", "-"]
 _Legacy: TypeAlias = Literal["1.13", "1.21", "1.25", "2.1", False]
 _LegacyNoStyle: TypeAlias = Literal["1.21", "1.25", "2.1", False]
-_ReprFunc: TypeAlias = Callable[[NDArray[Any]], str]
+_ReprFunc: TypeAlias = Callable[[_nt.Array], str]
 
 @type_check_only
 class _FormatDict(TypedDict, total=False):
@@ -104,7 +105,7 @@ def printoptions(
 # public numpy export
 @overload  # no style
 def array2string(
-    a: NDArray[Any],
+    a: _nt.Array,
     max_line_width: int | None = None,
     precision: SupportsIndex | None = None,
     suppress_small: bool | None = None,
@@ -122,7 +123,7 @@ def array2string(
 ) -> str: ...
 @overload  # style=<given> (positional), legacy="1.13"
 def array2string(
-    a: NDArray[Any],
+    a: _nt.Array,
     max_line_width: int | None,
     precision: SupportsIndex | None,
     suppress_small: bool | None,
@@ -140,7 +141,7 @@ def array2string(
 ) -> str: ...
 @overload  # style=<given> (keyword), legacy="1.13"
 def array2string(
-    a: NDArray[Any],
+    a: _nt.Array,
     max_line_width: int | None = None,
     precision: SupportsIndex | None = None,
     suppress_small: bool | None = None,
@@ -159,7 +160,7 @@ def array2string(
 @overload  # style=<given> (positional), legacy!="1.13"
 @deprecated("'style' argument is deprecated and no longer functional except in 1.13 'legacy' mode")
 def array2string(
-    a: NDArray[Any],
+    a: _nt.Array,
     max_line_width: int | None,
     precision: SupportsIndex | None,
     suppress_small: bool | None,
@@ -178,7 +179,7 @@ def array2string(
 @overload  # style=<given> (keyword), legacy="1.13"
 @deprecated("'style' argument is deprecated and no longer functional except in 1.13 'legacy' mode")
 def array2string(
-    a: NDArray[Any],
+    a: _nt.Array,
     max_line_width: int | None = None,
     precision: SupportsIndex | None = None,
     suppress_small: bool | None = None,
@@ -222,15 +223,15 @@ def format_float_positional(
 
 # public numpy export
 def array_repr(
-    arr: NDArray[Any],
+    arr: _nt.Array,
     max_line_width: int | None = None,
     precision: SupportsIndex | None = None,
     suppress_small: bool | None = None,
 ) -> str: ...
 
 # public numpy export
 def array_str(
-    a: NDArray[Any],
+    a: _nt.Array,
     max_line_width: int | None = None,
     precision: SupportsIndex | None = None,
     suppress_small: bool | None = None,