feat(typing): allow new_1d to take ndarrays

Author: alexfikl

doc/conf.py

@@ -45,6 +45,7 @@
     "np.dtype": "class:numpy.dtype",
     "np.ndarray": "class:numpy.ndarray",
     "np.floating": "class:numpy.floating",
+    "np.generic[Any]": "class:numpy.generic",
     # pytools typing
     "ObjectArray1D": "obj:pytools.obj_array.ObjectArray1D",
     "ReadableBuffer": "data:pytools.ReadableBuffer",

pytools/obj_array.py

@@ -7,6 +7,7 @@
 .. autoclass:: T_co
 .. autoclass:: ResultT
 .. autoclass:: ShapeT
+.. autoclass:: NumpyTypeT
 
 .. autoclass:: ObjectArray
 .. autoclass:: ObjectArray0D
@@ -99,6 +100,7 @@
 
 ResultT = TypeVar("ResultT")
 ShapeT = TypeVar("ShapeT", bound=tuple[int, ...])
+NumpyTypeT = TypeVar("NumpyTypeT", bound="np.generic[Any]")
 
 
 class _ObjectArrayMetaclass(type):
@@ -340,7 +342,37 @@ def from_numpy(
     return cast("ObjectArray[ShapeT, T_co]", cast("object", ary))
 
 
-def new_1d(res_list: Sequence[T_co]) -> ObjectArray1D[T_co]:
+@overload
+def new_1d(  # pyright: ignore[reportOverlappingOverload]
+        res_list: np.ndarray[tuple[int], np.dtype[NumpyTypeT]]
+    ) -> ObjectArray1D[NumpyTypeT]: ...
+
+@overload
+def new_1d(
+        res_list: np.ndarray[tuple[int, int], np.dtype[NumpyTypeT]]
+    ) -> ObjectArray1D[np.ndarray[tuple[int], np.dtype[NumpyTypeT]]]: ...
+
+@overload
+def new_1d(
+        res_list: np.ndarray[tuple[int, int, int], np.dtype[NumpyTypeT]]
+    ) -> ObjectArray1D[np.ndarray[tuple[int, int], np.dtype[NumpyTypeT]]]: ...
+
+@overload
+def new_1d(
+        res_list: np.ndarray[tuple[int, ...], np.dtype[NumpyTypeT]]
+    ) -> ObjectArray1D[np.ndarray[tuple[int, ...], np.dtype[NumpyTypeT]]]: ...
+
+@overload
+def new_1d(res_list: Sequence[T_co]) -> ObjectArray1D[T_co]: ...
+
+
+def new_1d(
+        res_list: (
+            Sequence[T_co]
+            | np.ndarray[tuple[int, ...], np.dtype[NumpyTypeT]]
+            )
+    ) -> (ObjectArray1D[T_co]
+          | ObjectArray1D[np.ndarray[tuple[int, ...], np.dtype[NumpyTypeT]]]):
     """Create a one-dimensional object array from *res_list*.
     This differs from ``numpy.array(res_list, dtype=object)``
     by whether it tries to determine its shape by descending