Merge branch 'main' into topk

* main:
  Fix first, last again (#381)
  Fix docs build (#382)
  Optimize for-loop merging of cohorts. (#378)
  Add cohorts snapshot tests with syrupy (#379)
  Stricter tolerance in property tests (#377)
  Better gridlines
  Update containment image with gridlines, higher dpi

Author: dcherian

.github/workflows/ci-additional.yaml

@@ -128,7 +128,8 @@ jobs:
 
       - name: Run mypy
         run: |
-          python -m mypy --install-types --non-interactive --cobertura-xml-report mypy_report
+          mkdir .mypy_cache
+          python -m mypy --install-types --non-interactive --cache-dir=.mypy_cache/ --cobertura-xml-report mypy_report
 
       - name: Upload mypy coverage to Codecov
         uses: codecov/[email protected]

.github/workflows/ci.yaml

@@ -16,7 +16,7 @@ concurrency:
 
 jobs:
   test:
-    name: Test (${{ matrix.python-version }}, ${{ matrix.os }})
+    name: Test (${{matrix.env}}, ${{ matrix.python-version }}, ${{ matrix.os }})
     runs-on: ${{ matrix.os }}
     defaults:
       run:
@@ -25,10 +25,18 @@ jobs:
       fail-fast: false
       matrix:
         os: ["ubuntu-latest"]
+        env: ["environment"]
         python-version: ["3.9", "3.12"]
         include:
           - os: "windows-latest"
+            env: "environment"
             python-version: "3.12"
+          - os: "ubuntu-latest"
+            env: "no-dask" # "no-xarray", "no-numba"
+            python-version: "3.12"
+          - os: "ubuntu-latest"
+            env: "minimal-requirements"
+            python-version: "3.9"
     steps:
       - uses: actions/checkout@v4
         with:
@@ -39,7 +47,7 @@ jobs:
       - name: Set up conda environment
         uses: mamba-org/setup-micromamba@v1
         with:
-          environment-file: ci/environment.yml
+          environment-file: ci/${{ matrix.env }}.yml
           environment-name: flox-tests
           init-shell: bash
           cache-environment: true
@@ -81,48 +89,6 @@ jobs:
           path: .hypothesis/
           key: cache-hypothesis-${{ runner.os }}-${{ matrix.python-version }}-${{ github.run_id }}
 
-  optional-deps:
-    name: ${{ matrix.env }}
-    runs-on: "ubuntu-latest"
-    defaults:
-      run:
-        shell: bash -l {0}
-    strategy:
-      fail-fast: false
-      matrix:
-        python-version: ["3.12"]
-        env: ["no-dask"] # "no-xarray", "no-numba"
-        include:
-          - env: "minimal-requirements"
-            python-version: "3.9"
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          fetch-depth: 0 # Fetch all history for all branches and tags.
-      - name: Set up conda environment
-        uses: mamba-org/setup-micromamba@v1
-        with:
-          environment-file: ci/${{ matrix.env }}.yml
-          environment-name: flox-tests
-          init-shell: bash
-          cache-environment: true
-          create-args: |
-            python=${{ matrix.python-version }}
-      - name: Install flox
-        run: |
-          python -m pip install --no-deps -e .
-      - name: Run tests
-        run: |
-          python -m pytest -n auto --cov=./ --cov-report=xml
-      - name: Upload code coverage to Codecov
-        uses: codecov/[email protected]
-        with:
-          file: ./coverage.xml
-          flags: unittests
-          env_vars: RUNNER_OS
-          name: codecov-umbrella
-          fail_ci_if_error: false
-
   xarray-groupby:
     name: xarray-groupby
     runs-on: ubuntu-latest

asv_bench/__init__.py

@@ -47,6 +47,9 @@ def time_find_group_cohorts(self):
         except AttributeError:
             pass
 
+    def track_num_cohorts(self):
+        return len(self.chunks_cohorts())
+
     def time_graph_construct(self):
         flox.groupby_reduce(self.array, self.by, func="sum", axis=self.axis)
 
@@ -60,10 +63,11 @@ def track_num_tasks_optimized(self):
     def track_num_layers(self):
         return len(self.result.dask.layers)
 
+    track_num_cohorts.unit = "cohorts"  # type: ignore[attr-defined] # Lazy
     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
-    for f in [track_num_tasks, track_num_tasks_optimized, track_num_layers]:
+    for f in [track_num_tasks, track_num_tasks_optimized, track_num_layers, track_num_cohorts]:
         f.repeat = 1  # type: ignore[attr-defined] # Lazy
         f.rounds = 1  # type: ignore[attr-defined] # Lazy
         f.number = 1  # type: ignore[attr-defined] # Lazy

asv_bench/benchmarks/cohorts.py

@@ -16,7 +16,7 @@ dependencies:
   - myst-parser
   - myst-nb
   - sphinx
-  - furo
+  - furo>=2024.08
   - ipykernel
   - jupyter
   - sphinx-codeautolink

ci/docs.yml

@@ -18,6 +18,7 @@ dependencies:
   - pytest-cov
   - pytest-pretty
   - pytest-xdist
+  - syrupy
   - xarray
   - pre-commit
   - numpy_groupies>=0.9.19

ci/environment.yml

@@ -9,6 +9,7 @@ dependencies:
   - pytest-cov
   - pytest-pretty
   - pytest-xdist
+  - syrupy
   - numpy==1.22
   - scipy==1.9.0
   - numpy_groupies==0.9.19

ci/minimal-requirements.yml

@@ -13,6 +13,7 @@ dependencies:
   - pytest-cov
   - pytest-pretty
   - pytest-xdist
+  - syrupy
   - xarray
   - numpydoc
   - pre-commit

ci/no-dask.yml

@@ -18,6 +18,7 @@ dependencies:
   - pytest-cov
   - pytest-pretty
   - pytest-xdist
+  - syrupy
   - xarray
   - pre-commit
   - numpy_groupies>=0.9.19

ci/no-numba.yml

@@ -3,6 +3,7 @@ channels:
   - conda-forge
 dependencies:
   - codecov
+  - syrupy
   - pandas
   - numpy>=1.22
   - scipy
@@ -11,6 +12,7 @@ dependencies:
   - pytest-cov
   - pytest-pretty
   - pytest-xdist
+  - syrupy
   - dask-core
   - numpydoc
   - pre-commit