From c0928dbca4368a037a6d277dd174eed9cd44d6ca Mon Sep 17 00:00:00 2001 From: jorenham Date: Tue, 4 Mar 2025 16:15:14 +0100 Subject: [PATCH] =?UTF-8?q?=F0=9F=8F=B7=EF=B8=8F=20fix=20stubtest=20errors?= =?UTF-8?q?=20in=20`numpy.lib.=5Fnanfunctions=5Fimpl`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .mypyignore-todo | 1 - src/numpy-stubs/lib/_function_base_impl.pyi | 3 +- src/numpy-stubs/lib/_nanfunctions_impl.pyi | 261 +++++++++++++++++++- 3 files changed, 261 insertions(+), 4 deletions(-) diff --git a/.mypyignore-todo b/.mypyignore-todo index 1d0488ba..c83f3a40 100644 --- a/.mypyignore-todo +++ b/.mypyignore-todo @@ -8,7 +8,6 @@ numpy(\..+)?\.floating.as_integer_ratio numpy(\..+)?\.complexfloating.__hash__ numpy(\..+)?\.complexfloating.__complex__ -numpy(\.lib\._nanfunctions_impl|\.matlib)?\.nan(median|percentile) numpy(\.lib\._twodim_base_impl)?\.tri(l|u) numpy\._?core(\._multiarray_umath|\.multiarray)\.error diff --git a/src/numpy-stubs/lib/_function_base_impl.pyi b/src/numpy-stubs/lib/_function_base_impl.pyi index 05ebb7a1..5c8704ff 100644 --- a/src/numpy-stubs/lib/_function_base_impl.pyi +++ b/src/numpy-stubs/lib/_function_base_impl.pyi @@ -700,7 +700,7 @@ def sinc(x: CoFloating_1nd) -> Array[np.floating]: ... @overload def sinc(x: ToComplex_1nd) -> Array[np.complexfloating]: ... -# +# keep in sync with `lib._nanfunctions_impl.nanmedian` @overload def median( a: CoFloating_nd, @@ -759,6 +759,7 @@ def median( keepdims: bool = False, ) -> _ArrayT: ... +# keep in sync with `lib._nanfunctions_impl.nanpercentile` # TODO(jorenham): deprecate interpolation # TODO(jorenham): deprecate only allow weights if method="inverted_cdf" @overload diff --git a/src/numpy-stubs/lib/_nanfunctions_impl.pyi b/src/numpy-stubs/lib/_nanfunctions_impl.pyi index 4a8225a4..f615a518 100644 --- a/src/numpy-stubs/lib/_nanfunctions_impl.pyi +++ b/src/numpy-stubs/lib/_nanfunctions_impl.pyi @@ -1,4 +1,20 @@ -# NOTE: In reality these functions are not aliases but distinct functions with identical signatures. +from typing import Any, Literal as L, overload +from typing_extensions import TypeVar + +import numpy as np +from _numtype import ( + Array, + CoComplex_nd, + CoDateTime_nd, + CoFloating_0d, + CoFloating_1nd, + CoFloating_nd, + CoTimeDelta_nd, + ToComplex_nd, + ToDateTime_nd, + ToObject_nd, + ToTimeDelta_nd, +) from numpy._core.fromnumeric import ( amax as nanmax, amin as nanmin, @@ -12,7 +28,10 @@ from numpy._core.fromnumeric import ( sum as nansum, var as nanvar, ) -from numpy.lib._function_base_impl import median as nanmedian, percentile as nanpercentile, quantile as nanquantile +from numpy._globals import _NoValueType +from numpy._typing import _ShapeLike + +from ._function_base_impl import _PercentileMethod __all__ = [ "nanargmax", @@ -30,3 +49,241 @@ __all__ = [ "nansum", "nanvar", ] + +### + +_ArrayT = TypeVar("_ArrayT", bound=Array) + +### + +# keep in sync with `lib._function_base_impl.median` +@overload +def nanmedian( + a: CoFloating_nd, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + keepdims: _NoValueType | L[False] = ..., +) -> np.floating: ... +@overload +def nanmedian( + a: ToComplex_nd, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + keepdims: _NoValueType | L[False] = ..., +) -> np.complexfloating: ... +@overload +def nanmedian( + a: ToTimeDelta_nd, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + keepdims: _NoValueType | L[False] = ..., +) -> np.timedelta64: ... +@overload +def nanmedian( + a: ToObject_nd, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + keepdims: _NoValueType | L[False] = ..., +) -> Any: ... +@overload +def nanmedian( + a: CoComplex_nd | CoTimeDelta_nd | ToObject_nd, + axis: _ShapeLike | None = None, + out: None = None, + overwrite_input: bool = False, + keepdims: _NoValueType | bool = ..., +) -> Any: ... +@overload +def nanmedian( + a: CoComplex_nd | CoTimeDelta_nd | ToObject_nd, + axis: _ShapeLike | None, + out: _ArrayT, + overwrite_input: bool = False, + keepdims: _NoValueType | bool = ..., +) -> _ArrayT: ... +@overload +def nanmedian( + a: CoComplex_nd | CoTimeDelta_nd | ToObject_nd, + axis: _ShapeLike | None = None, + *, + out: _ArrayT, + overwrite_input: bool = False, + keepdims: _NoValueType | bool = ..., +) -> _ArrayT: ... + +# keep in sync with `lib._function_base_impl.percentile` +@overload +def nanpercentile( + a: CoFloating_nd, + q: CoFloating_0d, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | L[False] = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> np.floating: ... +@overload +def nanpercentile( + a: CoFloating_nd, + q: CoFloating_1nd, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | L[False] = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> Array[np.floating]: ... +@overload +def nanpercentile( + a: ToComplex_nd, + q: CoFloating_0d, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | L[False] = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> np.complexfloating: ... +@overload +def nanpercentile( + a: ToComplex_nd, + q: CoFloating_1nd, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | L[False] = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> Array[np.complexfloating]: ... +@overload +def nanpercentile( + a: ToTimeDelta_nd, + q: CoFloating_0d, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | L[False] = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> np.timedelta64: ... +@overload +def nanpercentile( + a: ToTimeDelta_nd, + q: CoFloating_1nd, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | L[False] = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> Array[np.timedelta64]: ... +@overload +def nanpercentile( + a: ToDateTime_nd, + q: CoFloating_0d, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | L[False] = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> np.datetime64: ... +@overload +def nanpercentile( + a: ToDateTime_nd, + q: CoFloating_1nd, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | L[False] = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> Array[np.datetime64]: ... +@overload +def nanpercentile( + a: ToObject_nd, + q: CoFloating_0d, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | L[False] = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> Any: ... +@overload +def nanpercentile( + a: ToObject_nd, + q: CoFloating_1nd, + axis: None = None, + out: None = None, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | L[False] = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> Array[np.object_]: ... +@overload +def nanpercentile( + a: CoComplex_nd | CoDateTime_nd | ToObject_nd, + q: CoFloating_1nd, + axis: _ShapeLike | None = None, + out: None = None, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | bool = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> Any: ... +@overload +def nanpercentile( + a: CoComplex_nd | CoDateTime_nd | ToObject_nd, + q: CoFloating_1nd, + axis: _ShapeLike | None = None, + *, + out: _ArrayT, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | bool = ..., + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> _ArrayT: ... +@overload +def nanpercentile( + a: CoComplex_nd | CoDateTime_nd | ToObject_nd, + q: CoFloating_1nd, + axis: _ShapeLike | None, + out: _ArrayT, + overwrite_input: bool = False, + method: _PercentileMethod = "linear", + keepdims: _NoValueType | bool = ..., + *, + weights: CoFloating_1nd | None = None, + interpolation: None = None, +) -> _ArrayT: ... + +nanquantile = nanpercentile