Add type-hints to tests and misc

adaptive/notebook_integration.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import datetime
 import importlib
@@ -76,7 +78,7 @@ def ensure_plotly():
         raise RuntimeError("plotly is not installed; plotting is disabled.")
 
 
-def in_ipynb():
+def in_ipynb() -> bool:
     try:
         # If we are running in IPython, then `get_ipython()` is always a global
         return get_ipython().__class__.__name__ == "ZMQInteractiveShell"

adaptive/tests/test_learner1d.py

@@ -277,8 +277,8 @@ def f_vec(x, offset=0.123214):
     def assert_equal_dicts(d1, d2):
         xs1, ys1 = zip(*sorted(d1.items()))
         xs2, ys2 = zip(*sorted(d2.items()))
-        ys1 = np.array(ys1, dtype=np.float)
-        ys2 = np.array(ys2, dtype=np.float)
+        ys1 = np.array(ys1, dtype=np.float64)
+        ys2 = np.array(ys2, dtype=np.float64)
         np.testing.assert_almost_equal(xs1, xs2)
         np.testing.assert_almost_equal(ys1, ys2)
 

adaptive/types.py

@@ -11,3 +11,4 @@
 Float: TypeAlias = Union[float, np.float_]
 Int: TypeAlias = Union[int, np.int_]
 Real: TypeAlias = Union[Float, Int]
+Bool: TypeAlias = Union[bool, np.bool_]

adaptive/utils.py

@@ -1,24 +1,27 @@
+from __future__ import annotations
+
 import abc
 import functools
 import gzip
 import inspect
 import os
 import pickle
 import warnings
-from contextlib import contextmanager
+from contextlib import _GeneratorContextManager, contextmanager
 from itertools import product
+from typing import Any, Callable, Mapping, Sequence
 
 import cloudpickle
 
 
-def named_product(**items):
+def named_product(**items: Mapping[str, Sequence[Any]]):
     names = items.keys()
     vals = items.values()
     return [dict(zip(names, res)) for res in product(*vals)]
 
 
 @contextmanager
-def restore(*learners):
+def restore(*learners) -> _GeneratorContextManager:
     states = [learner.__getstate__() for learner in learners]
     try:
         yield
@@ -27,7 +30,7 @@ def restore(*learners):
             learner.__setstate__(state)
 
 
-def cache_latest(f):
+def cache_latest(f: Callable) -> Callable:
     """Cache the latest return value of the function and add it
     as 'self._cache[f.__name__]'."""
 
@@ -42,7 +45,7 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def save(fname, data, compress=True):
+def save(fname: str, data: Any, compress: bool = True) -> bool:
     fname = os.path.expanduser(fname)
     dirname = os.path.dirname(fname)
     if dirname:
@@ -71,14 +74,14 @@ def save(fname, data, compress=True):
     return True
 
 
-def load(fname, compress=True):
+def load(fname: str, compress: bool = True) -> Any:
     fname = os.path.expanduser(fname)
     _open = gzip.open if compress else open
     with _open(fname, "rb") as f:
         return cloudpickle.load(f)
 
 
-def copy_docstring_from(other):
+def copy_docstring_from(other: Callable) -> Callable:
     def decorator(method):
         return functools.wraps(other)(method)