ENH: Improve the annotations of `np.core._internal`

1 files changed, 32 insertions, 13 deletions

diff --git a/numpy/core/_internal.pyi b/numpy/core/_internal.pyi
index 1b3889e51..c24fc2ca3 100644
--- a/numpy/core/_internal.pyi
+++ b/numpy/core/_internal.pyi
@@ -1,18 +1,37 @@
-from typing import Any
+from typing import Any, TypeVar, Type, overload, Optional, Generic
+import ctypes as ct
 
-# TODO: add better annotations when ctypes is stubbed out
+from numpy import ndarray
 
-class _ctypes:
+_CastT = TypeVar("_CastT", bound=ct._CanCastTo)  # Copied from `ctypes.cast`
+_CT = TypeVar("_CT", bound=ct._CData)
+_PT = TypeVar("_PT", bound=Optional[int])
+
+# TODO: Let the likes of `shape_as` and `strides_as` return `None`
+# for 0D arrays once we've got shape-support
+
+class _ctypes(Generic[_PT]):
+    @overload
+    def __new__(cls, array: ndarray[Any, Any], ptr: None = ...) -> _ctypes[None]: ...
+    @overload
+    def __new__(cls, array: ndarray[Any, Any], ptr: _PT) -> _ctypes[_PT]: ...
+
+    # NOTE: In practice `shape` and `strides` return one of the concrete
+    # platform dependant c_int64-based array-types (`c_int`, `c_long` or
+    # `c_longlong`)
     @property
-    def data(self) -> int: ...
+    def data(self) -> _PT: ...
     @property
-    def shape(self) -> Any: ...
+    def shape(self) -> ct.Array[ct.c_int64]: ...
     @property
-    def strides(self) -> Any: ...
-    def data_as(self, obj: Any) -> Any: ...
-    def shape_as(self, obj: Any) -> Any: ...
-    def strides_as(self, obj: Any) -> Any: ...
-    def get_data(self) -> int: ...
-    def get_shape(self) -> Any: ...
-    def get_strides(self) -> Any: ...
-    def get_as_parameter(self) -> Any: ...
+    def strides(self) -> ct.Array[ct.c_int64]: ...
+    @property
+    def _as_parameter_(self) -> ct.c_void_p: ...
+
+    def data_as(self, obj: Type[_CastT]) -> _CastT: ...
+    def shape_as(self, obj: Type[_CT]) -> ct.Array[_CT]: ...
+    def strides_as(self, obj: Type[_CT]) -> ct.Array[_CT]: ...
+    def get_data(self) -> _PT: ...
+    def get_shape(self) -> ct.Array[ct.c_int64]: ...
+    def get_strides(self) -> ct.Array[ct.c_int64]: ...
+    def get_as_parameter(self) -> ct.c_void_p: ...