Merge pull request #155 from jenshnielsen/typecheck_numpy_stubs

Use numpy stups with mypy

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()

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
 

test_requirements.txt

@@ -3,3 +3,4 @@ pytest
 pytest-qt
 mypy
 PyQt5-stubs
+git+https://github.com/numpy/numpy-stubs.git