Merge pull request #19179 from BvB93/testing

ENH: Add annotations for `np.testing`

6 files changed, 444 insertions, 107 deletions

diff --git a/numpy/testing/__init__.py b/numpy/testing/__init__.py
index e1f87621f..064a74a16 100644
--- a/numpy/testing/__init__.py
+++ b/numpy/testing/__init__.py
@@ -8,6 +8,7 @@ away.
 from unittest import TestCase
 
 from ._private.utils import *
+from ._private.utils import _assert_valid_refcount, _gen_alignment_data
 from ._private import decorators as dec
 from ._private.nosetester import (
     run_module_suite, NoseTester as Tester
diff --git a/numpy/testing/__init__.pyi b/numpy/testing/__init__.pyi
index 395626f6b..955dae862 100644
--- a/numpy/testing/__init__.pyi
+++ b/numpy/testing/__init__.pyi
@@ -1,113 +1,53 @@
-import sys
-import warnings
-from typing import Any, List, ClassVar, Tuple, Set
-
-if sys.version_info >= (3, 8):
-    from typing import Final
-else:
-    from typing_extensions import Final
+from typing import List
 
 from unittest import (
     TestCase as TestCase,
 )
 
-from unittest.case import (
+from numpy.testing._private.utils import (
+    assert_equal as assert_equal,
+    assert_almost_equal as assert_almost_equal,
+    assert_approx_equal as assert_approx_equal,
+    assert_array_equal as assert_array_equal,
+    assert_array_less as assert_array_less,
+    assert_string_equal as assert_string_equal,
+    assert_array_almost_equal as assert_array_almost_equal,
+    assert_raises as assert_raises,
+    build_err_msg as build_err_msg,
+    decorate_methods as decorate_methods,
+    jiffies as jiffies,
+    memusage as memusage,
+    print_assert_equal as print_assert_equal,
+    raises as raises,
+    rundocs as rundocs,
+    runstring as runstring,
+    verbose as verbose,
+    measure as measure,
+    assert_ as assert_,
+    assert_array_almost_equal_nulp as assert_array_almost_equal_nulp,
+    assert_raises_regex as assert_raises_regex,
+    assert_array_max_ulp as assert_array_max_ulp,
+    assert_warns as assert_warns,
+    assert_no_warnings as assert_no_warnings,
+    assert_allclose as assert_allclose,
+    IgnoreException as IgnoreException,
+    clear_and_catch_warnings as clear_and_catch_warnings,
     SkipTest as SkipTest,
+    KnownFailureException as KnownFailureException,
+    temppath as temppath,
+    tempdir as tempdir,
+    IS_PYPY as IS_PYPY,
+    HAS_REFCOUNT as HAS_REFCOUNT,
+    suppress_warnings as suppress_warnings,
+    assert_array_compare as assert_array_compare,
+    assert_no_gc_cycles as assert_no_gc_cycles,
+    break_cycles as break_cycles,
+    HAS_LAPACK64 as HAS_LAPACK64,
 )
 
 __all__: List[str]
 
-def run_module_suite(file_to_run=..., argv=...): ...
-
-class KnownFailureException(Exception): ...
-class IgnoreException(Exception): ...
-
-class clear_and_catch_warnings(warnings.catch_warnings):
-    class_modules: ClassVar[Tuple[str, ...]]
-    modules: Set[str]
-    def __init__(self, record=..., modules=...): ...
-    def __enter__(self): ...
-    def __exit__(self, *exc_info): ...
-
-class suppress_warnings:
-    log: List[warnings.WarningMessage]
-    def __init__(self, forwarding_rule=...): ...
-    def filter(self, category=..., message=..., module=...): ...
-    def record(self, category=..., message=..., module=...): ...
-    def __enter__(self): ...
-    def __exit__(self, *exc_info): ...
-    def __call__(self, func): ...
-
-verbose: int
-IS_PYPY: Final[bool]
-HAS_REFCOUNT: Final[bool]
-HAS_LAPACK64: Final[bool]
-
-def assert_(val, msg=...): ...
-def memusage(processName=..., instance=...): ...
-def jiffies(_proc_pid_stat=..., _load_time=...): ...
-def build_err_msg(
-    arrays,
-    err_msg,
-    header=...,
-    verbose=...,
-    names=...,
-    precision=...,
-): ...
-def assert_equal(actual, desired, err_msg=..., verbose=...): ...
-def print_assert_equal(test_string, actual, desired): ...
-def assert_almost_equal(
-    actual,
-    desired,
-    decimal=...,
-    err_msg=...,
-    verbose=...,
-): ...
-def assert_approx_equal(
-    actual,
-    desired,
-    significant=...,
-    err_msg=...,
-    verbose=...,
-): ...
-def assert_array_compare(
-    comparison,
-    x,
-    y,
-    err_msg=...,
-    verbose=...,
-    header=...,
-    precision=...,
-    equal_nan=...,
-    equal_inf=...,
-): ...
-def assert_array_equal(x, y, err_msg=..., verbose=...): ...
-def assert_array_almost_equal(x, y, decimal=..., err_msg=..., verbose=...): ...
-def assert_array_less(x, y, err_msg=..., verbose=...): ...
-def runstring(astr, dict): ...
-def assert_string_equal(actual, desired): ...
-def rundocs(filename=..., raise_on_error=...): ...
-def raises(*args): ...
-def assert_raises(*args, **kwargs): ...
-def assert_raises_regex(exception_class, expected_regexp, *args, **kwargs): ...
-def decorate_methods(cls, decorator, testmatch=...): ...
-def measure(code_str, times=..., label=...): ...
-def assert_allclose(
-    actual,
-    desired,
-    rtol=...,
-    atol=...,
-    equal_nan=...,
-    err_msg=...,
-    verbose=...,
-): ...
-def assert_array_almost_equal_nulp(x, y, nulp=...): ...
-def assert_array_max_ulp(a, b, maxulp=..., dtype=...): ...
-def assert_warns(warning_class, *args, **kwargs): ...
-def assert_no_warnings(*args, **kwargs): ...
-def tempdir(*args, **kwargs): ...
-def temppath(*args, **kwargs): ...
-def assert_no_gc_cycles(*args, **kwargs): ...
-def break_cycles(): ...
-def _assert_valid_refcount(op): ...
-def _gen_alignment_data(dtype=..., type=..., max_size=...): ...
+def run_module_suite(
+    file_to_run: None | str = ...,
+    argv: None | List[str] = ...,
+) -> None: ...
diff --git a/numpy/testing/_private/utils.py b/numpy/testing/_private/utils.py
index 393fedc27..487aa0b4c 100644
--- a/numpy/testing/_private/utils.py
+++ b/numpy/testing/_private/utils.py
@@ -35,8 +35,7 @@ __all__ = [
         'assert_allclose', 'IgnoreException', 'clear_and_catch_warnings',
         'SkipTest', 'KnownFailureException', 'temppath', 'tempdir', 'IS_PYPY',
         'HAS_REFCOUNT', 'suppress_warnings', 'assert_array_compare',
-        '_assert_valid_refcount', '_gen_alignment_data', 'assert_no_gc_cycles',
-        'break_cycles', 'HAS_LAPACK64'
+        'assert_no_gc_cycles', 'break_cycles', 'HAS_LAPACK64'
         ]
 
 
@@ -2518,4 +2517,3 @@ def _no_tracing(func):
             finally:
                 sys.settrace(original_trace)
         return wrapper
-
diff --git a/numpy/testing/_private/utils.pyi b/numpy/testing/_private/utils.pyi
new file mode 100644
index 000000000..29915309f
--- /dev/null
+++ b/numpy/testing/_private/utils.pyi
@@ -0,0 +1,396 @@
+import os
+import sys
+import ast
+import types
+import warnings
+import unittest
+import contextlib
+from typing import (
+    Any,
+    AnyStr,
+    Callable,
+    ClassVar,
+    Dict,
+    Iterable,
+    List,
+    NoReturn,
+    overload,
+    Pattern,
+    Sequence,
+    Set,
+    Tuple,
+    Type,
+    type_check_only,
+    TypeVar,
+    Union,
+)
+
+from numpy import generic, dtype, number, object_, bool_, _FloatValue
+from numpy.typing import (
+    NDArray,
+    ArrayLike,
+    DTypeLike,
+    _ArrayLikeNumber_co,
+    _ArrayLikeObject_co,
+    _ArrayLikeTD64_co,
+    _ArrayLikeDT64_co,
+)
+
+from unittest.case import (
+    SkipTest as SkipTest,
+)
+
+if sys.version_info >= (3, 8):
+    from typing import Final, SupportsIndex, Literal as L
+else:
+    from typing_extensions import Final, SupportsIndex, Literal as L
+
+_T = TypeVar("_T")
+_ET = TypeVar("_ET", bound=BaseException)
+_FT = TypeVar("_FT", bound=Callable[..., Any])
+
+# Must return a bool or an ndarray/generic type
+# that is supported by `np.logical_and.reduce`
+_ComparisonFunc = Callable[
+    [NDArray[Any], NDArray[Any]],
+    Union[
+        bool,
+        bool_,
+        number[Any],
+        NDArray[Union[bool_, number[Any], object_]],
+    ],
+]
+
+__all__: List[str]
+
+class KnownFailureException(Exception): ...
+class IgnoreException(Exception): ...
+
+class clear_and_catch_warnings(warnings.catch_warnings):
+    class_modules: ClassVar[Tuple[types.ModuleType, ...]]
+    modules: Set[types.ModuleType]
+    @overload
+    def __new__(
+        cls,
+        record: L[False] = ...,
+        modules: Iterable[types.ModuleType] = ...,
+    ) -> _clear_and_catch_warnings_without_records: ...
+    @overload
+    def __new__(
+        cls,
+        record: L[True],
+        modules: Iterable[types.ModuleType] = ...,
+    ) -> _clear_and_catch_warnings_with_records: ...
+    @overload
+    def __new__(
+        cls,
+        record: bool,
+        modules: Iterable[types.ModuleType] = ...,
+    ) -> clear_and_catch_warnings: ...
+    def __enter__(self) -> None | List[warnings.WarningMessage]: ...
+    def __exit__(
+        self,
+        __exc_type: None | Type[BaseException] = ...,
+        __exc_val: None | BaseException = ...,
+        __exc_tb: None | types.TracebackType = ...,
+    ) -> None: ...
+
+# Type-check only `clear_and_catch_warnings` subclasses for both values of the
+# `record` parameter. Copied from the stdlib `warnings` stubs.
+
+@type_check_only
+class _clear_and_catch_warnings_with_records(clear_and_catch_warnings):
+    def __enter__(self) -> List[warnings.WarningMessage]: ...
+
+@type_check_only
+class _clear_and_catch_warnings_without_records(clear_and_catch_warnings):
+    def __enter__(self) -> None: ...
+
+class suppress_warnings:
+    log: List[warnings.WarningMessage]
+    def __init__(
+        self,
+        forwarding_rule: L["always", "module", "once", "location"] = ...,
+    ) -> None: ...
+    def filter(
+        self,
+        category: Type[Warning] = ...,
+        message: str = ...,
+        module: None | types.ModuleType = ...,
+    ) -> None: ...
+    def record(
+        self,
+        category: Type[Warning] = ...,
+        message: str = ...,
+        module: None | types.ModuleType = ...,
+    ) -> List[warnings.WarningMessage]: ...
+    def __enter__(self: _T) -> _T: ...
+    def __exit__(
+        self,
+        __exc_type: None | Type[BaseException] = ...,
+        __exc_val: None | BaseException = ...,
+        __exc_tb: None | types.TracebackType = ...,
+    ) -> None: ...
+    def __call__(self, func: _FT) -> _FT: ...
+
+verbose: int
+IS_PYPY: Final[bool]
+HAS_REFCOUNT: Final[bool]
+HAS_LAPACK64: Final[bool]
+
+def assert_(val: object, msg: str | Callable[[], str] = ...) -> None: ...
+
+# Contrary to runtime we can't do `os.name` checks while type checking,
+# only `sys.platform` checks
+if sys.platform == "win32" or sys.platform == "cygwin":
+    def memusage(processName: str = ..., instance: int = ...) -> int: ...
+elif sys.platform == "linux":
+    def memusage(_proc_pid_stat: str | bytes | os.PathLike[Any] = ...) -> None | int: ...
+else:
+    def memusage() -> NoReturn: ...
+
+if sys.platform == "linux":
+    def jiffies(
+        _proc_pid_stat: str | bytes | os.PathLike[Any] = ...,
+        _load_time: List[float] = ...,
+    ) -> int: ...
+else:
+    def jiffies(_load_time: List[float] = ...) -> int: ...
+
+def build_err_msg(
+    arrays: Iterable[object],
+    err_msg: str,
+    header: str = ...,
+    verbose: bool = ...,
+    names: Sequence[str] = ...,
+    precision: None | SupportsIndex = ...,
+) -> str: ...
+
+def assert_equal(
+    actual: object,
+    desired: object,
+    err_msg: str = ...,
+    verbose: bool = ...,
+) -> None: ...
+
+def print_assert_equal(
+    test_string: str,
+    actual: object,
+    desired: object,
+) -> None: ...
+
+def assert_almost_equal(
+    actual: _ArrayLikeNumber_co | _ArrayLikeObject_co,
+    desired: _ArrayLikeNumber_co | _ArrayLikeObject_co,
+    decimal: int = ...,
+    err_msg: str = ...,
+    verbose: bool = ...,
+) -> None: ...
+
+# Anything that can be coerced into `builtins.float`
+def assert_approx_equal(
+    actual: _FloatValue,
+    desired: _FloatValue,
+    significant: int = ...,
+    err_msg: str = ...,
+    verbose: bool = ...,
+) -> None: ...
+
+def assert_array_compare(
+    comparison: _ComparisonFunc,
+    x: ArrayLike,
+    y: ArrayLike,
+    err_msg: str = ...,
+    verbose: bool = ...,
+    header: str = ...,
+    precision: SupportsIndex = ...,
+    equal_nan: bool = ...,
+    equal_inf: bool = ...,
+) -> None: ...
+
+def assert_array_equal(
+    x: ArrayLike,
+    y: ArrayLike,
+    err_msg: str = ...,
+    verbose: bool = ...,
+) -> None: ...
+
+def assert_array_almost_equal(
+    x: _ArrayLikeNumber_co | _ArrayLikeObject_co,
+    y: _ArrayLikeNumber_co | _ArrayLikeObject_co,
+    decimal: float = ...,
+    err_msg: str = ...,
+    verbose: bool = ...,
+) -> None: ...
+
+@overload
+def assert_array_less(
+    x: _ArrayLikeNumber_co | _ArrayLikeObject_co,
+    y: _ArrayLikeNumber_co | _ArrayLikeObject_co,
+    err_msg: str = ...,
+    verbose: bool = ...,
+) -> None: ...
+@overload
+def assert_array_less(
+    x: _ArrayLikeTD64_co,
+    y: _ArrayLikeTD64_co,
+    err_msg: str = ...,
+    verbose: bool = ...,
+) -> None: ...
+@overload
+def assert_array_less(
+    x: _ArrayLikeDT64_co,
+    y: _ArrayLikeDT64_co,
+    err_msg: str = ...,
+    verbose: bool = ...,
+) -> None: ...
+
+def runstring(
+    astr: str | bytes | types.CodeType,
+    dict: None | Dict[str, Any],
+) -> Any: ...
+
+def assert_string_equal(actual: str, desired: str) -> None: ...
+
+def rundocs(
+    filename: None | str | os.PathLike[str] = ...,
+    raise_on_error: bool = ...,
+) -> None: ...
+
+def raises(*args: Type[BaseException]) -> Callable[[_FT], _FT]: ...
+
+@overload
+def assert_raises(  # type: ignore
+    __expected_exception: Type[BaseException] | Tuple[Type[BaseException], ...],
+    __callable: Callable[..., Any],
+    *args: Any,
+    **kwargs: Any,
+) -> None: ...
+@overload
+def assert_raises(
+    expected_exception: Type[_ET] | Tuple[Type[_ET], ...],
+    *,
+    msg: None | str = ...,
+) -> unittest.case._AssertRaisesContext[_ET]: ...
+
+@overload
+def assert_raises_regex(
+    __expected_exception: Type[BaseException] | Tuple[Type[BaseException], ...],
+    __expected_regex: str | bytes | Pattern[Any],
+    __callable: Callable[..., Any],
+    *args: Any,
+    **kwargs: Any,
+) -> None: ...
+@overload
+def assert_raises_regex(
+    expected_exception: Type[_ET] | Tuple[Type[_ET], ...],
+    expected_regex: str | bytes | Pattern[Any],
+    *,
+    msg: None | str = ...,
+) -> unittest.case._AssertRaisesContext[_ET]: ...
+
+def decorate_methods(
+    cls: Type[Any],
+    decorator: Callable[[Callable[..., Any]], Any],
+    testmatch: None | str | bytes | Pattern[Any] = ...,
+) -> None: ...
+
+def measure(
+    code_str: str | bytes | ast.mod | ast.AST,
+    times: int = ...,
+    label: None | str = ...,
+) -> float: ...
+
+@overload
+def assert_allclose(
+    actual: _ArrayLikeNumber_co | _ArrayLikeObject_co,
+    desired: _ArrayLikeNumber_co | _ArrayLikeObject_co,
+    rtol: float = ...,
+    atol: float = ...,
+    equal_nan: bool = ...,
+    err_msg: str = ...,
+    verbose: bool = ...,
+) -> None: ...
+@overload
+def assert_allclose(
+    actual: _ArrayLikeTD64_co,
+    desired: _ArrayLikeTD64_co,
+    rtol: float = ...,
+    atol: float = ...,
+    equal_nan: bool = ...,
+    err_msg: str = ...,
+    verbose: bool = ...,
+) -> None: ...
+
+def assert_array_almost_equal_nulp(
+    x: _ArrayLikeNumber_co,
+    y: _ArrayLikeNumber_co,
+    nulp: float = ...,
+) -> None: ...
+
+def assert_array_max_ulp(
+    a: _ArrayLikeNumber_co,
+    b: _ArrayLikeNumber_co,
+    maxulp: float = ...,
+    dtype: DTypeLike = ...,
+) -> NDArray[Any]: ...
+
+@overload
+def assert_warns(
+    warning_class: Type[Warning],
+) -> contextlib._GeneratorContextManager[None]: ...
+@overload
+def assert_warns(
+    __warning_class: Type[Warning],
+    __func: Callable[..., _T],
+    *args: Any,
+    **kwargs: Any,
+) -> _T: ...
+
+@overload
+def assert_no_warnings() -> contextlib._GeneratorContextManager[None]: ...
+@overload
+def assert_no_warnings(
+    __func: Callable[..., _T],
+    *args: Any,
+    **kwargs: Any,
+) -> _T: ...
+
+@overload
+def tempdir(
+    suffix: None = ...,
+    prefix: None = ...,
+    dir: None = ...,
+) -> contextlib._GeneratorContextManager[str]: ...
+@overload
+def tempdir(
+    suffix: None | AnyStr = ...,
+    prefix: None | AnyStr = ...,
+    dir: None | AnyStr | os.PathLike[AnyStr] = ...,
+) -> contextlib._GeneratorContextManager[AnyStr]: ...
+
+@overload
+def temppath(
+    suffix: None = ...,
+    prefix: None = ...,
+    dir: None = ...,
+    text: bool = ...,
+) -> contextlib._GeneratorContextManager[str]: ...
+@overload
+def temppath(
+    suffix: None | AnyStr = ...,
+    prefix: None | AnyStr = ...,
+    dir: None | AnyStr | os.PathLike[AnyStr] = ...,
+    text: bool = ...,
+) -> contextlib._GeneratorContextManager[AnyStr]: ...
+
+@overload
+def assert_no_gc_cycles() -> contextlib._GeneratorContextManager[None]: ...
+@overload
+def assert_no_gc_cycles(
+    __func: Callable[..., Any],
+    *args: Any,
+    **kwargs: Any,
+) -> None: ...
+
+def break_cycles() -> None: ...
diff --git a/numpy/testing/setup.py b/numpy/testing/setup.py
index 7652a94a2..6f203e872 100755
--- a/numpy/testing/setup.py
+++ b/numpy/testing/setup.py
@@ -7,6 +7,7 @@ def configuration(parent_package='',top_path=None):
     config.add_subpackage('_private')
     config.add_subpackage('tests')
     config.add_data_files('*.pyi')
+    config.add_data_files('_private/*.pyi')
     return config
 
 if __name__ == '__main__':
diff --git a/numpy/testing/utils.py b/numpy/testing/utils.py
index 753258c13..20a883304 100644
--- a/numpy/testing/utils.py
+++ b/numpy/testing/utils.py
@@ -12,6 +12,7 @@ warnings.warn("Importing from numpy.testing.utils is deprecated "
               DeprecationWarning, stacklevel=2)
 
 from ._private.utils import *
+from ._private.utils import _assert_valid_refcount, _gen_alignment_data
 
 __all__ = [
         'assert_equal', 'assert_almost_equal', 'assert_approx_equal',
@@ -24,5 +25,5 @@ __all__ = [
         'assert_allclose', 'IgnoreException', 'clear_and_catch_warnings',
         'SkipTest', 'KnownFailureException', 'temppath', 'tempdir', 'IS_PYPY',
         'HAS_REFCOUNT', 'suppress_warnings', 'assert_array_compare',
-        '_assert_valid_refcount', '_gen_alignment_data', 'assert_no_gc_cycles'
+        'assert_no_gc_cycles'
         ]