add type annotations for adaptive/tests/test_cquad.py

Author: basnijholt

adaptive/tests/test_cquad.py

@@ -1,28 +1,34 @@
 from functools import partial
 from operator import attrgetter
+from typing import Callable, List, Set, Union
 
 import numpy as np
 import pytest
 
-from adaptive.learner import IntegratorLearner
 from adaptive.learner.integrator_coeffs import ns
-from adaptive.learner.integrator_learner import DivergentIntegralError
+from adaptive.learner.integrator_learner import (
+    DivergentIntegralError,
+    IntegratorLearner,
+    _Interval,
+)
 
 from .algorithm_4 import DivergentIntegralError as A4DivergentIntegralError
 from .algorithm_4 import algorithm_4, f0, f7, f21, f24, f63, fdiv
 
 eps = np.spacing(1)
 
 
-def run_integrator_learner(f, a, b, tol, n):
+def run_integrator_learner(
+    f: Union[partial, Callable], a: int, b: int, tol: float, n: int
+) -> IntegratorLearner:
     learner = IntegratorLearner(f, bounds=(a, b), tol=tol)
     for _ in range(n):
         points, _ = learner.ask(1)
         learner.tell_many(points, map(learner.function, points))
     return learner
 
 
-def equal_ival(ival, other, *, verbose=False):
+def equal_ival(ival: _Interval, other: _Interval, *, verbose=False) -> bool:
     """Note: Implementing __eq__ breaks SortedContainers in some way."""
     if ival.depth_complete is None:
         if verbose:
@@ -42,7 +48,9 @@ def equal_ival(ival, other, *, verbose=False):
     return all(same_slots)
 
 
-def equal_ivals(ivals, other, *, verbose=False):
+def equal_ivals(
+    ivals: Set[_Interval], other: List[_Interval], *, verbose=False
+) -> bool:
     """Note: `other` is a list of ivals."""
     if len(ivals) != len(other):
         if verbose:
@@ -56,7 +64,7 @@ def equal_ivals(ivals, other, *, verbose=False):
     )
 
 
-def same_ivals(f, a, b, tol):
+def same_ivals(f: Callable, a: int, b: int, tol: float) -> bool:
     igral, err, n, ivals = algorithm_4(f, a, b, tol)
 
     learner = run_integrator_learner(f, a, b, tol, n)
@@ -71,15 +79,15 @@ def same_ivals(f, a, b, tol):
 
 # XXX: This *should* pass (https://github.com/python-adaptive/adaptive/issues/55)
 @pytest.mark.xfail
-def test_that_gives_same_intervals_as_reference_implementation():
+def test_that_gives_same_intervals_as_reference_implementation() -> None:
     for i, args in enumerate(
         [[f0, 0, 3, 1e-5], [f7, 0, 1, 1e-6], [f21, 0, 1, 1e-3], [f24, 0, 3, 1e-3]]
     ):
         assert same_ivals(*args), f"Function {i}"
 
 
 @pytest.mark.xfail
-def test_machine_precision():
+def test_machine_precision() -> None:
     f, a, b, tol = [partial(f63, alpha=0.987654321, beta=0.45), 0, 1, 1e-10]
     igral, err, n, ivals = algorithm_4(f, a, b, tol)
 
@@ -92,7 +100,7 @@ def test_machine_precision():
     assert equal_ivals(learner.ivals, ivals, verbose=True)
 
 
-def test_machine_precision2():
+def test_machine_precision2() -> None:
     f, a, b, tol = [partial(f63, alpha=0.987654321, beta=0.45), 0, 1, 1e-10]
     igral, err, n, ivals = algorithm_4(f, a, b, tol)
 
@@ -102,7 +110,7 @@ def test_machine_precision2():
     np.testing.assert_almost_equal(err, learner.err)
 
 
-def test_divergence():
+def test_divergence() -> None:
     """This function should raise a DivergentIntegralError."""
     f, a, b, tol = fdiv, 0, 1, 1e-6
     with pytest.raises(A4DivergentIntegralError) as e:
@@ -114,22 +122,22 @@ def test_divergence():
         run_integrator_learner(f, a, b, tol, n)
 
 
-def test_choosing_and_adding_points_one_by_one():
+def test_choosing_and_adding_points_one_by_one() -> None:
     learner = IntegratorLearner(f24, bounds=(0, 3), tol=1e-10)
     for _ in range(1000):
         xs, _ = learner.ask(1)
         for x in xs:
             learner.tell(x, learner.function(x))
 
 
-def test_choosing_and_adding_multiple_points_at_once():
+def test_choosing_and_adding_multiple_points_at_once() -> None:
     learner = IntegratorLearner(f24, bounds=(0, 3), tol=1e-10)
     xs, _ = learner.ask(100)
     for x in xs:
         learner.tell(x, learner.function(x))
 
 
-def test_adding_points_and_skip_one_point():
+def test_adding_points_and_skip_one_point() -> None:
     learner = IntegratorLearner(f24, bounds=(0, 3), tol=1e-10)
     xs, _ = learner.ask(17)
     skip_x = xs[1]
@@ -160,7 +168,7 @@ def test_adding_points_and_skip_one_point():
 
 # XXX: This *should* pass (https://github.com/python-adaptive/adaptive/issues/55)
 @pytest.mark.xfail
-def test_tell_in_random_order(first_add_33=False):
+def test_tell_in_random_order(first_add_33: bool = False) -> None:
     from operator import attrgetter
     import random
 
@@ -219,11 +227,11 @@ def test_tell_in_random_order(first_add_33=False):
 
 # XXX: This *should* pass (https://github.com/python-adaptive/adaptive/issues/55)
 @pytest.mark.xfail
-def test_tell_in_random_order_first_add_33():
+def test_tell_in_random_order_first_add_33() -> None:
     test_tell_in_random_order(first_add_33=True)
 
 
-def test_approximating_intervals():
+def test_approximating_intervals() -> None:
     import random
 
     learner = IntegratorLearner(f24, bounds=(0, 3), tol=1e-10)
@@ -252,7 +260,7 @@ def test_removed_choose_mutiple_points_at_once():
     assert list(learner.approximating_intervals)[0] == learner.first_ival
 
 
-def test_removed_ask_one_by_one():
+def test_removed_ask_one_by_one() -> None:
     with pytest.raises(RuntimeError):
         # This test should raise because integrating np.exp should be done
         # after the 33th point