diff options
Diffstat (limited to 'numpy/core')
| -rw-r--r-- | numpy/core/_internal.pyi | 45 |
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: ... |