Type strategies

Author: dcherian

tests/strategies.py

@@ -1,3 +1,7 @@
+from __future__ import annotations
+
+from typing import Any, Callable
+
 import cftime
 import dask
 import hypothesis.extra.numpy as npst
@@ -6,6 +10,8 @@
 
 from . import ALL_FUNCS, SCIPY_STATS_FUNCS
 
+Chunks = tuple[tuple[int, ...], ...]
+
 
 def supported_dtypes() -> st.SearchStrategy[np.dtype]:
     return (
@@ -38,7 +44,6 @@ def supported_dtypes() -> st.SearchStrategy[np.dtype]:
     elements={"allow_subnormal": False}, shape=npst.array_shapes(), dtype=supported_dtypes()
 )
 
-
 calendars = st.sampled_from(
     [
         "standard",
@@ -55,7 +60,7 @@ def supported_dtypes() -> st.SearchStrategy[np.dtype]:
 
 
 @st.composite
-def units(draw, *, calendar: str):
+def units(draw, *, calendar: str) -> str:
     choices = ["days", "hours", "minutes", "seconds", "milliseconds", "microseconds"]
     if calendar == "360_day":
         choices += ["months"]
@@ -67,20 +72,36 @@ def units(draw, *, calendar: str):
     year, month, day = dt.year, dt.month, dt.day
     if calendar == "360_day":
         day = min(day, 30)
+    if calendar in ["360_day", "365_day", "noleap"] and month == 2 and day == 29:
+        day = 28
+
     return f"{time_units} since {year}-{month}-{day}"
 
 
 @st.composite
-def cftime_arrays(draw, *, shape, calendars=calendars, elements=None):
+def cftime_arrays(
+    draw: st.DrawFn,
+    *,
+    shape: tuple[int, ...],
+    calendars: st.SearchStrategy[str] = calendars,
+    elements: dict[str, Any] | None = None,
+) -> np.ndarray[Any, Any]:
     if elements is None:
-        elements = {"min_value": -10_000, "max_value": 10_000}
+        elements = {}
+    elements.setdefault("min_value", -10_000)
+    elements.setdefault("max_value", 10_000)
     cal = draw(calendars)
     values = draw(npst.arrays(dtype=np.int64, shape=shape, elements=elements))
     unit = draw(units(calendar=cal))
     return cftime.num2date(values, units=unit, calendar=cal)
 
 
-def by_arrays(shape, *, elements=None):
+def by_arrays(
+    shape: tuple[int, ...], *, elements: dict[str, Any] | None = None
+) -> st.SearchStrategy[np.ndarray[Any, Any]]:
+    if elements is None:
+        elements = {}
+    elements.setdefault("alphabet", st.characters(exclude_categories=("C",)))
     return st.one_of(
         npst.arrays(
             dtype=npst.integer_dtypes(endianness="=") | npst.unicode_string_dtypes(endianness="="),
@@ -92,7 +113,7 @@ def by_arrays(shape, *, elements=None):
 
 
 @st.composite
-def chunks(draw, *, shape: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
+def chunks(draw: st.DrawFn, *, shape: tuple[int, ...]) -> Chunks:
     chunks = []
     for size in shape:
         if size > 1:
@@ -107,7 +128,13 @@ def chunks(draw, *, shape: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
 
 
 @st.composite
-def chunked_arrays(draw, *, chunks=chunks, arrays=numeric_arrays, from_array=dask.array.from_array):
+def chunked_arrays(
+    draw: st.DrawFn,
+    *,
+    chunks: Callable[..., st.SearchStrategy[Chunks]] = chunks,
+    arrays=all_arrays,
+    from_array: Callable = dask.array.from_array,
+) -> dask.array.Array:
     array = draw(arrays)
     chunks = draw(chunks(shape=array.shape))
 

tests/test_properties.py

@@ -1,3 +1,5 @@
+from typing import Any, Callable
+
 import pandas as pd
 import pytest
 
@@ -14,7 +16,7 @@
 from flox.core import groupby_reduce, groupby_scan
 
 from . import assert_equal
-from .strategies import all_arrays, by_arrays, chunked_arrays, func_st, numeric_arrays
+from .strategies import by_arrays, chunked_arrays, func_st, numeric_arrays
 
 dask.config.set(scheduler="sync")
 
@@ -31,14 +33,14 @@ def bfill(array, axis, dtype=None):
     )[::-1]
 
 
-NUMPY_SCAN_FUNCS = {
+NUMPY_SCAN_FUNCS: dict[str, Callable] = {
     "nancumsum": np.nancumsum,
     "ffill": ffill,
     "bfill": bfill,
 }  # "cumsum": np.cumsum,
 
 
-def not_overflowing_array(array) -> bool:
+def not_overflowing_array(array: np.ndarray[Any, Any]) -> bool:
     if array.dtype.kind == "f":
         info = np.finfo(array.dtype)
     elif array.dtype.kind in ["i", "u"]:
@@ -51,9 +53,8 @@ def not_overflowing_array(array) -> bool:
     return result
 
 
-# TODO: migrate to by_arrays but with constant value
 @given(data=st.data(), array=numeric_arrays, func=func_st)
-def test_groupby_reduce(data, array, func):
+def test_groupby_reduce(data, array, func: str) -> None:
     # overflow behaviour differs between bincount and sum (for example)
     assume(not_overflowing_array(array))
     # TODO: fix var for complex numbers upstream
@@ -71,14 +72,14 @@ def test_groupby_reduce(data, array, func):
                 "min_size": 1,
                 "max_size": 1,
             },
-            shape=array.shape[-1],
+            shape=(array.shape[-1],),
         )
     )
     assert len(np.unique(by)) == 1
     kwargs = {"q": 0.8} if "quantile" in func else {}
-    flox_kwargs = {}
+    flox_kwargs: dict[str, Any] = {}
     with np.errstate(invalid="ignore", divide="ignore"):
-        actual, _ = groupby_reduce(
+        actual, *_ = groupby_reduce(
             array, by, func=func, axis=axis, engine="numpy", **flox_kwargs, finalize_kwargs=kwargs
         )
 
@@ -112,10 +113,10 @@ def test_groupby_reduce(data, array, func):
 
 @given(
     data=st.data(),
-    array=chunked_arrays(),
+    array=chunked_arrays(arrays=numeric_arrays),
     func=st.sampled_from(tuple(NUMPY_SCAN_FUNCS)),
 )
-def test_scans(data, array, func):
+def test_scans(data, array: dask.array.Array, func: str) -> None:
     assume(not_overflowing_array(np.asarray(array)))
 
     by = data.draw(by_arrays(shape=(array.shape[-1],)))
@@ -148,7 +149,7 @@ def test_scans(data, array, func):
 
 
 @given(data=st.data(), array=chunked_arrays())
-def test_ffill_bfill_reverse(data, array):
+def test_ffill_bfill_reverse(data, array: dask.array.Array) -> None:
     # TODO: test NaT and timedelta, datetime
     assume(not_overflowing_array(np.asarray(array)))
     by = data.draw(by_arrays(shape=(array.shape[-1],)))
@@ -168,10 +169,10 @@ def reverse(arr):
 
 @given(
     data=st.data(),
-    array=chunked_arrays(arrays=all_arrays),
+    array=chunked_arrays(),
     func=st.sampled_from(["first", "last", "nanfirst", "nanlast"]),
 )
-def test_first_last(data, array, func):
+def test_first_last(data, array: dask.array.Array, func: str) -> None:
     by = data.draw(by_arrays(shape=(array.shape[-1],)))
 
     INVERSES = {"first": "last", "last": "first", "nanfirst": "nanlast", "nanlast": "nanfirst"}
@@ -183,8 +184,8 @@ def test_first_last(data, array, func):
         array = array.rechunk((*array.chunks[:-1], -1))
 
     for arr in [array, array.compute()]:
-        forward, fg = groupby_reduce(arr, by, func=func, engine="flox")
-        reverse, rg = groupby_reduce(arr[..., ::-1], by[..., ::-1], func=inverse, engine="flox")
+        forward, *fg = groupby_reduce(arr, by, func=func, engine="flox")
+        reverse, *rg = groupby_reduce(arr[..., ::-1], by[..., ::-1], func=inverse, engine="flox")
 
         assert forward.dtype == reverse.dtype
         assert forward.dtype == arr.dtype
@@ -196,6 +197,6 @@ def test_first_last(data, array, func):
         if mate in ["first", "last"]:
             array = array.rechunk((*array.chunks[:-1], -1))
 
-        first, _ = groupby_reduce(array, by, func=func, engine="flox")
-        second, _ = groupby_reduce(array, by, func=mate, engine="flox")
+        first, *_ = groupby_reduce(array, by, func=func, engine="flox")
+        second, *_ = groupby_reduce(array, by, func=mate, engine="flox")
         assert_equal(first, second)