From 70ff5aaa83041a6653d4283d74f68c5e582a9128 Mon Sep 17 00:00:00 2001 From: Charles Burkland Date: Mon, 11 Aug 2025 13:35:37 -0700 Subject: [PATCH] Makes function signatures as explicit as possible --- src/__init__.pyi | 44 +++++++++++++++++++++++++++++++------------- src/methods.c | 2 +- 2 files changed, 32 insertions(+), 14 deletions(-) diff --git a/src/__init__.pyi b/src/__init__.pyi index 888bb58c..347c256b 100644 --- a/src/__init__.pyi +++ b/src/__init__.pyi @@ -26,6 +26,16 @@ _TLabel = tp.Union[ tp.Tuple['_TLabel', ...], ] +_TVShape = tp.TypeVar("_TVShape") +_1D = tp.Tuple[tp.Any] +_2D = tp.Tuple[tp.Any, tp.Any] +_TVDtypeAny = tp.TypeVar('_TVDtypeAny', bound=np.dtype) +_TVNDArrayAny = tp.TypeVar("_TVNDArrayAny", np.ndarray) +_T1DArrayInt64 = np.ndarray[_1D, np.dtype[np.int64]] +_T1DArrayBool = np.ndarray[_1D, np.dtype[np.bool_]] +_T1DArrayObject = np.ndarray[_1D, np.dtype[np.object_]] +_T2DArrayBool = np.ndarray[_2D, np.dtype[np.bool_]] + class ErrorInitTypeBlocks(RuntimeError): def __init__(self, *args: tp.Any, **kwargs: tp.Any) -> None: ... def with_traceback(self, tb: Exception) -> Exception: ... @@ -119,7 +129,7 @@ class FrozenAutoMap: def __init__(self, labels: tp.Iterable[_TLabel] | np.ndarray = (), /,) -> None: ... def get(self, __key: _TLabel, /,) -> int: ... def keys(self) -> tp.Iterator[_TLabel]: ... - def items(self) -> tp.Iterator[tuple[_TLabel, int]]: ... + def items(self) -> tp.Iterator[tp.Tuple[_TLabel, int]]: ... def values(self) -> tp.Iterator[int]: ... def get_all(self, __key: list[_TLabel] | np.ndarray) -> np.ndarray: ... def get_any(self, __key: list[_TLabel] | np.ndarray) -> list[int]: ... @@ -186,26 +196,34 @@ def split_after_count( def count_iteration(__iterable: tp.Iterable) -> int: ... -def immutable_filter(__array: np.ndarray) -> np.ndarray: ... +def immutable_filter(__array: _TVNDArrayAny) -> _TVNDArrayAny: ... def mloc(__array: np.ndarray) -> int: ... def name_filter(__name: _TLabel) -> _TLabel: ... -def shape_filter(__array: np.ndarray) -> np.ndarray: ... -def column_2d_filter(__array: np.ndarray) -> np.ndarray: ... -def column_1d_filter(__array: np.ndarray) -> np.ndarray: ... -def row_1d_filter(__array: np.ndarray) -> np.ndarray: ... -def array_deepcopy(__array: np.ndarray, memo: tp.Optional[tp.Dict[int, tp.Any]]) -> np.ndarray: ... + +def shape_filter(__array: np.ndarray[_1D, _TVDtypeAny] | np.ndarray[_2D, _TVDtypeAny]) -> np.ndarray[_2D, _TVDtypeAny]: ... +def column_2d_filter(__array: np.ndarray[_1D, _TVDtypeAny] | np.ndarray[_2D, _TVDtypeAny]) -> np.ndarray[_2D, _TVDtypeAny]: ... +def column_1d_filter(__array: np.ndarray[_1D, _TVDtypeAny] | np.ndarray[_2D, _TVDtypeAny]) -> np.ndarray[_1D, _TVDtypeAny]: ... +def row_1d_filter(__array: np.ndarray[_1D, _TVDtypeAny] | np.ndarray[_2D, _TVDtypeAny]) -> np.ndarray[_1D, _TVDtypeAny]: ... + +def array_deepcopy(__array: _TVNDArrayAny, memo: tp.Optional[tp.Dict[int, tp.Any]]) -> _TVNDArrayAny: ... def resolve_dtype(__d1: np.dtype, __d2: np.dtype) -> np.dtype: ... def resolve_dtype_iter(__dtypes: tp.Iterable[np.dtype]) -> np.dtype: ... def isna_element(__value: tp.Any, include_none: bool = True) -> bool: ... def dtype_from_element(__value: tp.Optional[_TLabel]) -> np.dtype: ... -def get_new_indexers_and_screen(indexers: np.ndarray, positions: np.ndarray) -> tp.Tuple[np.ndarray, np.ndarray]: ... +def get_new_indexers_and_screen(indexers: _T1DArrayInt64, positions: _T1DArrayInt64) -> tp.Tuple[_T1DArrayInt64, _T1DArrayInt64]: ... -def first_true_1d(__array: np.ndarray, *, forward: bool) -> int: ... -def first_true_2d(__array: np.ndarray, *, forward: bool, axis: int) -> np.ndarray: ... -def nonzero_1d(__array: np.ndarray, /) -> np.ndarray: ... +def first_true_1d(__array: _T1DArrayBool, *, forward: bool) -> int: ... +def first_true_2d(__array: _T2DArrayBool, *, forward: bool, axis: int) -> _T1DArrayInt64: ... +def nonzero_1d(__array: _T1DArrayBool, /) -> _T1DArrayInt64: ... def is_objectable_dt64(__array: np.ndarray, /) -> bool: ... def is_objectable(__array: np.ndarray, /) -> bool: ... -def astype_array(__array: np.ndarray, __dtype: np.dtype | None, /) -> np.ndarray: ... + +@tp.overload +def astype_array(__array: np.ndarray[_TVShape, tp.Any], __dtype: _TVDtypeAny, /) -> np.ndarray[_TVShape, _TVDtypeAny]: ... +@tp.overload +def astype_array(__array: np.ndarray[_TVShape, tp.Any], __dtype: None, /) -> np.ndarray[_TVShape, np.int64]: ... +def astype_array(__array: np.ndarray[_TVShape, tp.Any], __dtype: _TVDtypeAny | None, /) -> np.ndarray[_TVShape, _TVDtypeAny] | np.ndarray[_TVShape, np.int64]: ... + def slice_to_ascending_slice(__slice: slice, __size: int) -> slice: ... -def array_to_tuple_array(__array: np.ndarray) -> np.ndarray: ... +def array_to_tuple_array(__array: np.ndarray) -> _T1DArrayObject: ... def array_to_tuple_iter(__array: np.ndarray) -> tp.Iterator[tp.Tuple[tp.Any, ...]]: ... \ No newline at end of file diff --git a/src/methods.c b/src/methods.c index d98f75cd..8514ff92 100644 --- a/src/methods.c +++ b/src/methods.c @@ -870,7 +870,7 @@ get_new_indexers_and_screen(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kw } if (PyArray_TYPE(indexers) != NPY_INT64) { - PyErr_SetString(PyExc_ValueError, "Array must be of type np.int64"); + PyErr_SetString(PyExc_ValueError, "indexers must be of type np.int64"); return NULL; }