Add type-hints to adaptive/utils.py

Author: basnijholt

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) -> None:
     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):
     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)