diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 01cd0c875c84..00b6ac729fc2 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -3679,6 +3679,8 @@ class memmap(ndarray[_ShapeType, _DType_co]): ) -> Any: ... def flush(self) -> None: ... +# TODO: Add a mypy plugin for managing functions whose output type is dependant +# on the literal value of some sort of signature (e.g. `einsum` and `vectorize`) class vectorize: pyfunc: Callable[..., Any] cache: bool @@ -3695,7 +3697,7 @@ class vectorize: cache: bool = ..., signature: None | str = ..., ) -> None: ... - def __call__(self, *args: Any, **kwargs: Any) -> NDArray[Any]: ... + def __call__(self, *args: Any, **kwargs: Any) -> Any: ... class poly1d: @property diff --git a/numpy/typing/tests/data/reveal/lib_function_base.pyi b/numpy/typing/tests/data/reveal/lib_function_base.pyi index c559eb295893..eebe9fbfdc70 100644 --- a/numpy/typing/tests/data/reveal/lib_function_base.pyi +++ b/numpy/typing/tests/data/reveal/lib_function_base.pyi @@ -26,7 +26,7 @@ reveal_type(vectorized_func.signature) # E: Union[None, builtins.str] reveal_type(vectorized_func.otypes) # E: Union[None, builtins.str] reveal_type(vectorized_func.excluded) # E: set[Union[builtins.int, builtins.str]] reveal_type(vectorized_func.__doc__) # E: Union[None, builtins.str] -reveal_type(vectorized_func([1])) # E: ndarray[Any, dtype[Any]] +reveal_type(vectorized_func([1])) # E: Any reveal_type(np.vectorize(int)) # E: vectorize reveal_type(np.vectorize( # E: vectorize int, otypes="i", doc="doc", excluded=(), cache=True, signature=None