typing

Author: dcherian

flox/aggregations.py

@@ -593,9 +593,6 @@ class Scan:
     mode: T_ScanBinaryOpMode = "apply_binary_op"
 
 
-HANDLED_FUNCTIONS = {}
-
-
 def concatenate(arrays: Sequence[AlignedArrays], axis=-1, out=None) -> AlignedArrays:
     group_idx = np.concatenate([a.group_idx for a in arrays], axis=axis)
     array = np.concatenate([a.array for a in arrays], axis=axis)
@@ -642,16 +639,16 @@ def __post_init__(self):
         assert (self.state is not None) or (self.result is not None)
 
 
-def scan_binary_op(
-    left_state: AlignedArrays, right_state: AlignedArrays, *, agg: Scan
-) -> AlignedArrays:
+def scan_binary_op(left_state: ScanState, right_state: ScanState, *, agg: Scan) -> ScanState:
     from .core import reindex_
 
     assert left_state.state is not None
     left = left_state.state
     right = right_state.result if right_state.result is not None else right_state.state
+    assert right is not None
 
     if agg.mode == "apply_binary_op":
+        assert agg.binary_op is not None
         # Implements groupby binary operation.
         reindexed = reindex_(
             left.array,
@@ -718,7 +715,7 @@ def scan_binary_op(
 # cumprod = Scan("cumprod", binary_op=np.multiply, preop="prod", scan="cumprod")
 
 
-AGGREGATIONS = {
+AGGREGATIONS: dict[str, Aggregation | Scan] = {
     "any": any_,
     "all": all_,
     "count": count,
@@ -769,7 +766,9 @@ def _initialize_aggregation(
         try:
             # TODO: need better interface
             # we set dtype, fillvalue on reduction later. so deepcopy now
-            agg = copy.deepcopy(AGGREGATIONS[func])
+            agg_ = copy.deepcopy(AGGREGATIONS[func])
+            assert isinstance(agg_, Aggregation)
+            agg = agg_
         except KeyError:
             raise NotImplementedError(f"Reduction {func!r} not implemented yet")
     elif isinstance(func, Aggregation):

flox/core.py

@@ -814,7 +814,7 @@ def factorize_(
                     idx = sorter[(idx,)]
                 idx[mask] = -1
             else:
-                idx, groups = pd.factorize(flat, sort=sort)  # type: ignore[arg-type]
+                idx, groups = pd.factorize(flat, sort=sort)
 
             found_groups.append(np.array(groups))
         factorized.append(idx.reshape(groupvar.shape))
@@ -2717,7 +2717,7 @@ def groupby_scan(
         agg.dtype = array.dtype
 
     if not has_dask:
-        (single_axis,) = axis_
+        (single_axis,) = axis_  # type: ignore[misc]
         final_state = chunk_scan(
             AlignedArrays(array=array, group_idx=by_), axis=single_axis, agg=agg, dtype=agg.dtype
         )
@@ -2765,6 +2765,7 @@ def _zip(group_idx: np.ndarray, array: np.ndarray) -> AlignedArrays:
 
 
 def extract_array(block: ScanState) -> np.ndarray:
+    assert block.result is not None
     return block.result.array
 
 
@@ -2787,7 +2788,7 @@ def dask_groupby_scan(array, by, axes: T_Axes, agg: Scan) -> DaskArray:
 
     scan_ = partial(chunk_scan, agg=agg)
     # dask tokenizing error workaround
-    scan_.__name__ = scan_.func.__name__
+    scan_.__name__ = scan_.func.__name__  # type: ignore[attr-defined]
 
     # 2. Run the scan
     accumulated = scan(

tests/test_core.py

@@ -1833,3 +1833,9 @@ def test_nanlen_string(dtype, engine):
 #     func="cumsum",
 #     axis=-1,
 # )
+
+# numpy_array, group_idx = (
+#     array([1.6777218e07, 1.0000000e00, 0.0000000e00], dtype=float32),
+#     array([0, 1, 1]),
+# )
+# groupby_scan(numpy_array, group_idx, axis=-1, func="nancumsum")