TYP: annotate strings.slice

jorenham · charris · commit 3ef41f72616c · 2025-05-27T12:53:37.000-06:00
diff --git a/numpy/_core/strings.pyi b/numpy/_core/strings.pyi
@@ -54,6 +54,7 @@ __all__ = [
     "translate",
     "upper",
     "zfill",
+    "slice",
 ]
 
 _StringDTypeArray: TypeAlias = np.ndarray[_AnyShape, np.dtypes.StringDType]
@@ -494,3 +495,17 @@ def translate(
     table: str,
     deletechars: str | None = None,
 ) -> _StringDTypeOrUnicodeArray: ...
+
+#
+@overload
+def slice(a: U_co, start: i_co | None = None, stop: i_co | None = None, step: i_co | None = None, /) -> NDArray[np.str_]: ...  # type: ignore[overload-overlap]
+@overload
+def slice(a: S_co, start: i_co | None = None, stop: i_co | None = None, step: i_co | None = None, /) -> NDArray[np.bytes_]: ...
+@overload
+def slice(
+    a: _StringDTypeSupportsArray, start: i_co | None = None, stop: i_co | None = None, step: i_co | None = None, /
+) -> _StringDTypeArray: ...
+@overload
+def slice(
+    a: T_co, start: i_co | None = None, stop: i_co | None = None, step: i_co | None = None, /
+) -> _StringDTypeOrUnicodeArray: ...
diff --git a/numpy/strings/__init__.pyi b/numpy/strings/__init__.pyi
@@ -36,6 +36,7 @@ from numpy._core.strings import (
     rjust,
     rpartition,
     rstrip,
+    slice,
     startswith,
     str_len,
     strip,
@@ -92,4 +93,5 @@ __all__ = [
     "decode",
     "encode",
     "translate",
+    "slice",
 ]
diff --git a/numpy/typing/tests/data/reveal/strings.pyi b/numpy/typing/tests/data/reveal/strings.pyi
@@ -190,3 +190,7 @@ assert_type(np.strings.str_len(AR_T), npt.NDArray[np.int_])
 assert_type(np.strings.translate(AR_U, ""), npt.NDArray[np.str_])
 assert_type(np.strings.translate(AR_S, ""), npt.NDArray[np.bytes_])
 assert_type(np.strings.translate(AR_T, ""), AR_T_alias)
+
+assert_type(np.strings.slice(AR_U, 1, 5, 2), npt.NDArray[np.str_])
+assert_type(np.strings.slice(AR_S, 1, 5, 2), npt.NDArray[np.bytes_])
+assert_type(np.strings.slice(AR_T, 1, 5, 2), AR_T_alias)
