From 0cd921d390c20988bf6013a5916ffb1ca17eb1f0 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sat, 21 Dec 2024 18:33:58 +0100 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20`=5Flib`:=20improved=20annotations?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../_lib/_elementwise_iterative_method.pyi | 103 +++++++++++++++++- scipy-stubs/_lib/_finite_differences.pyi | 5 + scipy-stubs/_lib/_gcutils.pyi | 6 + scipy-stubs/_lib/decorator.pyi | 2 + 4 files changed, 115 insertions(+), 1 deletion(-) diff --git a/scipy-stubs/_lib/_elementwise_iterative_method.pyi b/scipy-stubs/_lib/_elementwise_iterative_method.pyi index f93c5705..9e68e1c3 100644 --- a/scipy-stubs/_lib/_elementwise_iterative_method.pyi +++ b/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: ... diff --git a/scipy-stubs/_lib/_finite_differences.pyi b/scipy-stubs/_lib/_finite_differences.pyi index 50386b19..1c3d8fd8 100644 --- a/scipy-stubs/_lib/_finite_differences.pyi +++ b/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], diff --git a/scipy-stubs/_lib/_gcutils.pyi b/scipy-stubs/_lib/_gcutils.pyi index 1693d4bf..d0d1aed2 100644 --- a/scipy-stubs/_lib/_gcutils.pyi +++ b/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): ... diff --git a/scipy-stubs/_lib/decorator.pyi b/scipy-stubs/_lib/decorator.pyi index 500aaf34..76c24cdc 100644 --- a/scipy-stubs/_lib/decorator.pyi +++ b/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]]: ...