diff --git a/plottr/data/datadict_storage.py b/plottr/data/datadict_storage.py index 9f5ae32c..98a1eacc 100644 --- a/plottr/data/datadict_storage.py +++ b/plottr/data/datadict_storage.py @@ -17,7 +17,7 @@ import os import time from enum import Enum -from typing import Any, Union, Optional, Dict, Type +from typing import Any, Union, Optional, Dict, Type, Collection from types import TracebackType import numpy as np @@ -334,7 +334,7 @@ def datadict_from_hdf5(basepath: str, for k in keys: ds = grp[k] - entry = dict(values=np.array([]), ) + entry: Dict[str, Union[Collection[Any], np.ndarray]] = dict(values=np.array([]), ) if 'axes' in ds.attrs: entry['axes'] = deh5ify(ds.attrs['axes']).tolist() diff --git a/plottr/utils/num.py b/plottr/utils/num.py index bfda0f00..86f0b8a0 100644 --- a/plottr/utils/num.py +++ b/plottr/utils/num.py @@ -46,7 +46,7 @@ def largest_numtype(arr: np.ndarray, include_integers: bool = True) \ return None -def _are_close(a: np.ndarray, b: np.ndarray, rtol: float = 1e-8) -> np.ndarray: +def _are_close(a: np.ndarray, b: np.ndarray, rtol: float = 1e-8) -> Union[np.ndarray, np.bool_]: return np.isclose(a, b, rtol=rtol) @@ -98,7 +98,7 @@ def arrays_equal(a: np.ndarray, b: np.ndarray, else: return False - close = np.zeros(a.shape, dtype=bool) + close: Union[np.ndarray, np.bool_] = np.zeros(a.shape, dtype=bool) if a.dtype in FLOATTYPES and b.dtype in FLOATTYPES: close = _are_close(a, b, rtol=rtol) @@ -178,7 +178,7 @@ def _find_switches(arr: np.ndarray, def find_direction_period(vals: np.ndarray, ignore_last: bool = False) \ - -> Union[None, int]: + -> Optional[float]: """ Find the period with which the values in an array change direction. @@ -234,7 +234,7 @@ def guess_grid_from_sweep_direction(**axes: np.ndarray) \ :raises: `ValueError` for incorrect input """ periods_list = [] - sorting = [] + sorting: List[float] = [] names_list = [] size: Optional[int] = None diff --git a/test_requirements.txt b/test_requirements.txt index ce7657ec..138113c3 100644 --- a/test_requirements.txt +++ b/test_requirements.txt @@ -3,3 +3,4 @@ pytest pytest-qt mypy PyQt5-stubs +git+https://github.com/numpy/numpy-stubs.git \ No newline at end of file