diff options
Diffstat (limited to 'numpy/core/_internal.pyi')
| -rw-r--r-- | numpy/core/_internal.pyi | 38 |
1 files changed, 25 insertions, 13 deletions
diff --git a/numpy/core/_internal.pyi b/numpy/core/_internal.pyi index 1b3889e51..f4bfd770f 100644 --- a/numpy/core/_internal.pyi +++ b/numpy/core/_internal.pyi @@ -1,18 +1,30 @@ -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 +from numpy.ctypeslib import c_intp -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]: ... @property - def data(self) -> int: ... + def data(self) -> _PT: ... @property - def shape(self) -> Any: ... + def shape(self) -> ct.Array[c_intp]: ... @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[c_intp]: ... + @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]: ... |