mypy

Author: awf

src/gfloat/decode_ndarray.py

@@ -1,11 +1,13 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
 
+from types import ModuleType
 import numpy as np
+from .types import FormatInfo
 
-from .types import FloatClass, FloatValue, FormatInfo
 
-
-def decode_ndarray(fi: FormatInfo, codes: np.ndarray, np=np) -> np.ndarray:
+def decode_ndarray(
+    fi: FormatInfo, codes: np.ndarray, np: ModuleType = np
+) -> np.ndarray:
     r"""
     Vectorized version of :function:`decode_float`
 

test/test_decode.py

@@ -34,7 +34,7 @@ def dec(code: int) -> float:
 
 
 @pytest.mark.parametrize("method", methods)
-def test_spot_check_ocp_e5m2(method) -> None:
+def test_spot_check_ocp_e5m2(method: str) -> None:
     fi = format_info_ocp_e5m2
     dec = get_method(method, fi)
     fclass = lambda code: decode_float(fi, code).fclass
@@ -51,7 +51,7 @@ def test_spot_check_ocp_e5m2(method) -> None:
 
 
 @pytest.mark.parametrize("method", methods)
-def test_spot_check_ocp_e4m3(method) -> None:
+def test_spot_check_ocp_e4m3(method: str) -> None:
     fi = format_info_ocp_e4m3
     dec = get_method(method, fi)
     assert dec(0x40) == 2.0