✨ _lib: improved annotations (#364)

scipy-stubs/_lib/_elementwise_iterative_method.pyi

@@ -1 +1,102 @@
-# nothing to see here
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from types import ModuleType
+from typing import Any, Concatenate, Final, Literal, TypeAlias
+from typing_extensions import TypeVar
+
+import numpy as np
+import optype as op
+import optype.numpy as onp
+from ._util import _RichResult
+
+###
+
+_FloatT = TypeVar("_FloatT", bound=np.floating[Any], default=np.floating[Any])
+_ShapeT = TypeVar("_ShapeT", bound=onp.AtLeast1D, default=onp.AtLeast1D)
+_FuncRealT = TypeVar("_FuncRealT", bound=Callable[Concatenate[onp.ArrayND[np.float64], ...], object])
+_ModuleT = TypeVar("_ModuleT", bound=ModuleType, default=ModuleType)
+_WorkT = TypeVar("_WorkT", bound=Mapping[str, Any])
+_ResT = TypeVar("_ResT", bound=_RichResult, default=_RichResult)
+_ToShapeT = TypeVar("_ToShapeT", bound=op.CanIndex | tuple[op.CanIndex, ...], default=tuple[int, ...])
+
+_Ignored: TypeAlias = _ResT
+
+###
+
+_ESIGNERR: Final = -1
+_ECONVERR: Final = -2
+_EVALUEERR: Final = -3
+_ECALLBACK: Final = -4
+_EINPUTERR: Final = -5
+_ECONVERGED: Final = 0
+_EINPROGRESS: Final = 1
+
+# TODO: complex
+def _initialize(
+    func: _FuncRealT,
+    xs: Sequence[onp.ToFloat1D],
+    args: tuple[onp.ToFloat1D, ...],
+    complex_ok: Literal[False] = False,
+    preserve_shape: bool | None = None,
+    xp: _ModuleT | None = None,
+) -> tuple[
+    _FuncRealT,  # func
+    list[onp.Array1D[_FloatT]],  # xs
+    list[onp.Array1D[_FloatT]],  # fs
+    list[onp.Array1D[np.floating[Any]]],  # args
+    onp.AtLeast1D,  # shape
+    _FloatT,  # xfat
+    _ModuleT,  # xp
+]: ...
+
+# TODO: `_RichResult` subtype
+def _loop(
+    work: _ResT,
+    callback: Callable[[_ResT], _Ignored],
+    shape: Sequence[op.CanIndex],
+    maxiter: int,
+    func: Callable[[onp.Array[_ShapeT, _FloatT]], onp.ToComplexND],
+    args: tuple[onp.ArrayND[np.floating[Any]], ...],
+    dtype: np.inexact[Any],
+    pre_func_eval: Callable[[_ResT], onp.Array[_ShapeT, _FloatT]],
+    post_func_eval: Callable[[onp.Array[_ShapeT, _FloatT], onp.Array[_ShapeT, np.floating[Any]], _ResT], _Ignored],
+    check_termination: Callable[[_ResT], onp.Array[_ShapeT, np.bool_]],
+    post_termination_check: Callable[[_ResT], _Ignored],
+    customize_result: Callable[[_ResT, _ToShapeT], tuple[int, ...]],
+    res_work_pairs: Iterable[tuple[str, str]],
+    xp: ModuleType,
+    preserve_shape: bool | None = False,
+) -> _ResT: ...
+
+#
+def _check_termination(
+    work: _WorkT,
+    res: Mapping[str, onp.Array[_ShapeT, _FloatT]],
+    res_work_pairs: Iterable[tuple[str, str]],
+    active: onp.Array[_ShapeT, np.integer[Any]],
+    check_termination: Callable[[_WorkT], onp.Array[_ShapeT, np.bool_]],
+    preserve_shape: bool | None,
+    xp: ModuleType,
+) -> onp.Array1D[np.intp]: ...
+
+#
+def _update_active(
+    work: Mapping[str, onp.Array[_ShapeT, _FloatT]],
+    res: Mapping[str, onp.Array[_ShapeT, _FloatT]],
+    res_work_pairs: Iterable[tuple[str, str]],
+    active: onp.Array[_ShapeT, np.integer[Any]],
+    mask: onp.Array[_ShapeT, np.bool_] | None,
+    preserve_shape: bool | None,
+    xp: ModuleType,
+) -> None: ...
+
+#
+def _prepare_result(
+    work: Mapping[str, onp.Array[_ShapeT, _FloatT]],
+    res: _ResT,
+    res_work_pairs: Iterable[tuple[str, str]],
+    active: onp.Array[_ShapeT, np.integer[Any]],
+    shape: _ToShapeT,
+    customize_result: Callable[[_ResT, _ToShapeT], tuple[int, ...]],
+    preserve_shape: bool | None,
+    xp: ModuleType,
+) -> _ResT: ...

scipy-stubs/_lib/_finite_differences.pyi

@@ -7,6 +7,11 @@ import optype.numpy as onp
 
 _T = TypeVar("_T", bound=np.floating[Any] | onp.ArrayND[np.floating[Any]])
 
+###
+
+def _central_diff_weights(Np: int, ndiv: int = 1) -> onp.Array1D[np.float64]: ...
+
+#
 @overload
 def _derivative(
     func: Callable[Concatenate[float, ...], onp.ToFloat],

scipy-stubs/_lib/_gcutils.pyi

@@ -2,9 +2,15 @@ from contextlib import _GeneratorContextManager
 from collections.abc import Callable
 from typing import Final, ParamSpec, TypeVar
 
+__all__ = ["assert_deallocated", "gc_state", "set_gc_state"]
+
+###
+
 _T = TypeVar("_T")
 _Tss = ParamSpec("_Tss")
 
+###
+
 IS_PYPY: Final[bool] = ...
 
 class ReferenceError(AssertionError): ...

scipy-stubs/_lib/decorator.pyi

@@ -14,6 +14,8 @@ class ContextManager(_GeneratorContextManager[_T_co], Generic[_T_co]):
     def __init__(self, /, g: Callable[_Tss, Iterator[_T_co]], *a: _Tss.args, **k: _Tss.kwargs) -> None: ...
 
 def contextmanager(g: Callable[_Tss, Iterator[_T]], *a: _Tss.args, **k: _Tss.kwargs) -> Callable[[], ContextManager[_T]]: ...
+
+#
 def decorate(func: _FT, caller: Callable[Concatenate[_FT, _Tss], _T]) -> Callable[_Tss, _T]: ...
 @overload
 def decorator(caller: Callable[Concatenate[_FT, _Tss], _T], _func: None = None) -> Callable[[_FT], Callable[_Tss, _T]]: ...