diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/__init__.pyi | 2 | ||||
| -rw-r--r-- | numpy/core/_internal.pyi | 45 |
2 files changed, 33 insertions, 14 deletions
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 2d23f926d..f23dee9e2 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -1668,7 +1668,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]): @overload def __array__(self, __dtype: _DType) -> ndarray[Any, _DType]: ... @property - def ctypes(self) -> _ctypes: ... + def ctypes(self) -> _ctypes[int]: ... @property def shape(self) -> _Shape: ... @shape.setter 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: ... |