mypy

Author: awf

src/gfloat/round.py

@@ -1,4 +1,5 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+from typing import Callable
 
 import math
 
@@ -16,7 +17,7 @@ def _iseven(v: int) -> bool:
     return v & 0x1 == 0
 
 
-def _rnitp(x: float, pred) -> int:
+def _rnitp(x: float, pred: Callable) -> int:
     """Round to nearest integer, ties to predicate"""
     floored = math.floor(x)
     should_round_away = (x > floored + 0.5) | ((x == floored + 0.5) & ~pred(floored))

src/gfloat/round_ndarray.py

@@ -1,13 +1,13 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
 
-from typing import Optional, Tuple
+from typing import Callable, Optional, Tuple
 from .types import FormatInfo, RoundMode, Domain
 
 import numpy.typing as npt
 import array_api_compat
 
 
-def _ifloor(x, int_type):
+def _ifloor(x: npt.NDArray, int_type: npt.DTypeLike) -> npt.NDArray:
     xp = array_api_compat.array_namespace(x)
     floored = xp.floor(x)
     return xp.astype(floored, int_type)
@@ -21,7 +21,7 @@ def _iseven(v: npt.NDArray) -> npt.NDArray:
     return v & 0x1 == 0
 
 
-def _rnitp(x, pred, int_type):
+def _rnitp(x: npt.NDArray, pred: Callable[[npt.NDArray], npt.NDArray], int_type: npt.DTypeLike) -> npt.NDArray:
     """Round to nearest integer, ties to predicate"""
     xp = array_api_compat.array_namespace(x)
     floored = xp.floor(x)
@@ -31,12 +31,12 @@ def _rnitp(x, pred, int_type):
     return ifloored + xp.astype(should_round_away, int_type)
 
 
-def _rnite(x, int_type):
+def _rnite(x: npt.NDArray, int_type: npt.DTypeLike) -> npt.NDArray:
     """Round to nearest integer, ties to even"""
     return _rnitp(x, _iseven, int_type)
 
 
-def _rnito(x, int_type):
+def _rnito(x: npt.NDArray, int_type: npt.DTypeLike) -> npt.NDArray:
     """Round to nearest integer, ties to odd"""
     return _rnitp(x, _isodd, int_type)
 

test/test_round.py

@@ -5,6 +5,7 @@
 import array_api_compat
 import ml_dtypes
 import numpy as np
+import numpy.typing as npt
 import pytest
 
 from gfloat import RoundMode, decode_float, decode_ndarray, round_float, round_ndarray
@@ -14,7 +15,7 @@
 
 
 @pytest.mark.parametrize("int_type", [np.int64, np.int16])
-def test_rnito_rnite(int_type):
+def test_rnito_rnite(int_type : npt.DTypeLike) -> None:
 
     xp = array_api_compat.array_namespace(np.array(0.0))
     np.testing.assert_equal(_rnito(xp.array(3.5), int_type), 3.0)