From 40b300c9418b931da3cceaa9c799f55eab7d59a3 Mon Sep 17 00:00:00 2001
From: Illviljan <14371165+Illviljan@users.noreply.github.com>
Date: Sat, 8 Jul 2023 14:28:59 +0200
Subject: [PATCH] Fix some typing errors in asv_bench and tests (#253)

* Let mypy check asv_bench

* Update pyproject.toml

* Update pyproject.toml

* Update pyproject.toml

* Update pyproject.toml

* Update pyproject.toml

* Update pyproject.toml

* Update test_core.py

* Update test_core.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update test_core.py

* Update cohorts.py

* fix asv test

* split-reduce is old name

* Update test_core.py

* the exclude doesn't seem to work anyway

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 asv_bench/benchmarks/cohorts.py |  6 +++---
 asv_bench/benchmarks/combine.py |  2 +-
 pyproject.toml                  |  6 +++---
 tests/test_core.py              | 14 ++++++++------
 4 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/asv_bench/benchmarks/cohorts.py b/asv_bench/benchmarks/cohorts.py
index ecef95d22..dbfbe8cd5 100644
--- a/asv_bench/benchmarks/cohorts.py
+++ b/asv_bench/benchmarks/cohorts.py
@@ -42,9 +42,9 @@ def track_num_layers(self):
         )[0]
         return len(result.dask.layers)
 
-    track_num_tasks.unit = "tasks"
-    track_num_tasks_optimized.unit = "tasks"
-    track_num_layers.unit = "layers"
+    track_num_tasks.unit = "tasks"  # type: ignore[attr-defined] # Lazy
+    track_num_tasks_optimized.unit = "tasks"  # type: ignore[attr-defined] # Lazy
+    track_num_layers.unit = "layers"  # type: ignore[attr-defined] # Lazy
 
 
 class NWMMidwest(Cohorts):
diff --git a/asv_bench/benchmarks/combine.py b/asv_bench/benchmarks/combine.py
index dd3d7a178..27600685f 100644
--- a/asv_bench/benchmarks/combine.py
+++ b/asv_bench/benchmarks/combine.py
@@ -70,7 +70,7 @@ def construct_member(groups) -> dict[str, Any]:
         ]
         self.kwargs = {
             "agg": flox.aggregations._initialize_aggregation(
-                "sum", "float64", np.float64, 0, None, {}
+                "sum", "float64", np.float64, 0, 0, {}
             ),
             "axis": (3,),
         }
diff --git a/pyproject.toml b/pyproject.toml
index 360be8a27..fb27ee761 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -104,8 +104,7 @@ known-third-party = [
 
 [tool.mypy]
 allow_redefinition = true
-exclude = "properties|asv_bench|doc|tests|flycheck"
-files = "flox/*.py"
+files = "**/*.py"
 show_error_codes = true
 warn_unused_ignores = true
 
@@ -115,7 +114,8 @@ module=[
     "cftime",
     "dask.*",
     "importlib_metadata",
-    "numpy_groupies",
+    "numba",
+    "numpy_groupies.*",
     "matplotlib.*",
     "pandas",
     "setuptools",
diff --git a/tests/test_core.py b/tests/test_core.py
index 5c4db9248..83b823b07 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -78,7 +78,7 @@ def dask_array_ones(*args):
 )
 
 if TYPE_CHECKING:
-    from flox.core import T_Engine, T_ExpectedGroupsOpt, T_Func2
+    from flox.core import T_Agg, T_Engine, T_ExpectedGroupsOpt, T_Method
 
 
 def _get_array_func(func: str) -> Callable:
@@ -135,7 +135,7 @@ def test_alignment_error():
 )
 def test_groupby_reduce(
     engine: T_Engine,
-    func: T_Func2,
+    func: T_Agg,
     array: np.ndarray,
     by: np.ndarray,
     expected: list[float],
@@ -992,14 +992,14 @@ def test_map_reduce_blockwise_mixed() -> None:
         dask.array.from_array(data.values, chunks=365),
         t.dt.month,
         func="mean",
-        method="split-reduce",
+        method="map-reduce",
     )
     expected, _ = groupby_reduce(data, t.dt.month, func="mean")
     assert_equal(expected, actual)
 
 
 @requires_dask
-@pytest.mark.parametrize("method", ["split-reduce", "blockwise", "map-reduce", "cohorts"])
+@pytest.mark.parametrize("method", ["blockwise", "map-reduce", "cohorts"])
 def test_group_by_datetime(engine, method):
     kwargs = dict(
         func="mean",
@@ -1356,13 +1356,15 @@ def test_validate_reindex_map_reduce(
 
 
 def test_validate_reindex() -> None:
-    for method in ["map-reduce", "cohorts"]:
+    methods: list[T_Method] = ["map-reduce", "cohorts"]
+    for method in methods:
         with pytest.raises(NotImplementedError):
             _validate_reindex(
                 True, "argmax", method, expected_groups=None, any_by_dask=False, is_dask_array=True
             )
 
-    for method in ["blockwise", "cohorts"]:
+    methods: list[T_Method] = ["blockwise", "cohorts"]
+    for method in methods:
         with pytest.raises(ValueError):
             _validate_reindex(
                 True, "sum", method, expected_groups=None, any_by_dask=False, is_dask_array=True