More precisely type pipe methods (#10038)

chuckwondo · web-flow · commit 0caf09628011 · 2025-02-18T17:11:39.000-08:00
* Upgrade mypy to 1.15 Mypy 1.15 includes fix for <python/mypy#9031>, allowing several "type: ignore" comments to be removed. * Add type annotations to DataTree.pipe tests * More precisely type `pipe` methods. In addition, enhance mypy job configuration to support running it locally via `act`. Fixes #9997 * Pin mypy to 1.15 in CI * Revert mypy CI job changes * Add pytest-mypy-plugin and typestub packages * Add pytest-mypy-plugins to all conda env files * Remove dup pandas-stubs dep * Revert pre-commit config changes * Place mypy tests behind pytest mypy marker * Set default pytest numprocesses to 4 * Ignore pytest-mypy-plugins for min version check
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -67,6 +67,16 @@ jobs:
           - env: "flaky"
             python-version: "3.13"
             os: ubuntu-latest
+          # The mypy tests must be executed using only 1 process in order to guarantee
+          # predictable mypy output messages for comparison to expectations.
+          - env: "mypy"
+            python-version: "3.10"
+            numprocesses: 1
+            os: ubuntu-latest
+          - env: "mypy"
+            python-version: "3.13"
+            numprocesses: 1
+            os: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
         with:
@@ -88,6 +98,10 @@ jobs:
             then
               echo "CONDA_ENV_FILE=ci/requirements/environment.yml" >> $GITHUB_ENV
               echo "PYTEST_ADDOPTS=-m 'flaky or network' --run-flaky --run-network-tests -W default" >> $GITHUB_ENV
+            elif [[ "${{ matrix.env }}" == "mypy" ]] ;
+            then
+              echo "CONDA_ENV_FILE=ci/requirements/environment.yml" >> $GITHUB_ENV
+              echo "PYTEST_ADDOPTS=-n 1 -m 'mypy' --run-mypy -W default" >> $GITHUB_ENV
             else
               echo "CONDA_ENV_FILE=ci/requirements/${{ matrix.env }}.yml" >> $GITHUB_ENV
             fi
@@ -144,7 +158,7 @@ jobs:
           save-always: true
 
       - name: Run tests
-        run: python -m pytest -n 4
+        run: python -m pytest -n ${{ matrix.numprocesses || 4 }}
           --timeout 180
           --cov=xarray
           --cov-report=xml
diff --git a/ci/minimum_versions.py b/ci/minimum_versions.py
@@ -26,8 +26,9 @@
     "pytest",
     "pytest-cov",
     "pytest-env",
-    "pytest-xdist",
+    "pytest-mypy-plugins",
     "pytest-timeout",
+    "pytest-xdist",
     "hypothesis",
 ]
 
diff --git a/ci/requirements/all-but-dask.yml b/ci/requirements/all-but-dask.yml
@@ -30,8 +30,9 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-env
-  - pytest-xdist
+  - pytest-mypy-plugins
   - pytest-timeout
+  - pytest-xdist
   - rasterio
   - scipy
   - seaborn
diff --git a/ci/requirements/all-but-numba.yml b/ci/requirements/all-but-numba.yml
@@ -43,8 +43,9 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-env
-  - pytest-xdist
+  - pytest-mypy-plugins
   - pytest-timeout
+  - pytest-xdist
   - rasterio
   - scipy
   - seaborn
diff --git a/ci/requirements/bare-minimum.yml b/ci/requirements/bare-minimum.yml
@@ -9,8 +9,9 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-env
-  - pytest-xdist
+  - pytest-mypy-plugins
   - pytest-timeout
+  - pytest-xdist
   - numpy=1.24
   - packaging=23.1
   - pandas=2.1
diff --git a/ci/requirements/environment-3.14.yml b/ci/requirements/environment-3.14.yml
@@ -29,6 +29,7 @@ dependencies:
   - opt_einsum
   - packaging
   - pandas
+  - pandas-stubs
   # - pint>=0.22
   - pip
   - pooch
@@ -38,14 +39,25 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-env
-  - pytest-xdist
+  - pytest-mypy-plugins
   - pytest-timeout
+  - pytest-xdist
   - rasterio
   - scipy
   - seaborn
   # - sparse
   - toolz
+  - types-colorama
+  - types-docutils
+  - types-psutil
+  - types-Pygments
+  - types-python-dateutil
+  - types-pytz
+  - types-PyYAML
+  - types-setuptools
   - typing_extensions
   - zarr
   - pip:
       - jax # no way to get cpu-only jaxlib from conda if gpu is present
+      - types-defusedxml
+      - types-pexpect
diff --git a/ci/requirements/environment-windows-3.14.yml b/ci/requirements/environment-windows-3.14.yml
@@ -25,6 +25,7 @@ dependencies:
   - numpy
   - packaging
   - pandas
+  - pandas-stubs
   # - pint>=0.22
   - pip
   - pre-commit
@@ -33,12 +34,24 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-env
-  - pytest-xdist
+  - pytest-mypy-plugins
   - pytest-timeout
+  - pytest-xdist
   - rasterio
   - scipy
   - seaborn
   # - sparse
   - toolz
+  - types-colorama
+  - types-docutils
+  - types-psutil
+  - types-Pygments
+  - types-python-dateutil
+  - types-pytz
+  - types-PyYAML
+  - types-setuptools
   - typing_extensions
   - zarr
+  - pip:
+      - types-defusedxml
+      - types-pexpect
diff --git a/ci/requirements/environment-windows.yml b/ci/requirements/environment-windows.yml
@@ -25,6 +25,7 @@ dependencies:
   - numpy
   - packaging
   - pandas
+  - pandas-stubs
   # - pint>=0.22
   - pip
   - pre-commit
@@ -33,12 +34,24 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-env
-  - pytest-xdist
+  - pytest-mypy-plugins
   - pytest-timeout
+  - pytest-xdist
   - rasterio
   - scipy
   - seaborn
   - sparse
   - toolz
+  - types-colorama
+  - types-docutils
+  - types-psutil
+  - types-Pygments
+  - types-python-dateutil
+  - types-pytz
+  - types-PyYAML
+  - types-setuptools
   - typing_extensions
   - zarr
+  - pip:
+      - types-defusedxml
+      - types-pexpect
diff --git a/ci/requirements/environment.yml b/ci/requirements/environment.yml
@@ -29,6 +29,7 @@ dependencies:
   - opt_einsum
   - packaging
   - pandas
+  - pandas-stubs
   # - pint>=0.22
   - pip
   - pooch
@@ -39,14 +40,25 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-env
-  - pytest-xdist
+  - pytest-mypy-plugins
   - pytest-timeout
+  - pytest-xdist
   - rasterio
   - scipy
   - seaborn
   - sparse
   - toolz
+  - types-colorama
+  - types-docutils
+  - types-psutil
+  - types-Pygments
+  - types-python-dateutil
+  - types-pytz
+  - types-PyYAML
+  - types-setuptools
   - typing_extensions
   - zarr
   - pip:
       - jax # no way to get cpu-only jaxlib from conda if gpu is present
+      - types-defusedxml
+      - types-pexpect
diff --git a/ci/requirements/min-all-deps.yml b/ci/requirements/min-all-deps.yml
@@ -46,8 +46,9 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-env
-  - pytest-xdist
+  - pytest-mypy-plugins
   - pytest-timeout
+  - pytest-xdist
   - rasterio=1.3
   - scipy=1.11
   - seaborn=0.13
diff --git a/conftest.py b/conftest.py
@@ -3,14 +3,15 @@
 import pytest
 
 
-def pytest_addoption(parser):
+def pytest_addoption(parser: pytest.Parser):
     """Add command-line flags for pytest."""
     parser.addoption("--run-flaky", action="store_true", help="runs flaky tests")
     parser.addoption(
         "--run-network-tests",
         action="store_true",
         help="runs tests requiring a network connection",
     )
+    parser.addoption("--run-mypy", action="store_true", help="runs mypy tests")
 
 
 def pytest_runtest_setup(item):
@@ -21,6 +22,21 @@ def pytest_runtest_setup(item):
         pytest.skip(
             "set --run-network-tests to run test requiring an internet connection"
         )
+    if "mypy" in item.keywords and not item.config.getoption("--run-mypy"):
+        pytest.skip("set --run-mypy option to run mypy tests")
+
+
+# See https://docs.pytest.org/en/stable/example/markers.html#automatically-adding-markers-based-on-test-names
+def pytest_collection_modifyitems(items):
+    for item in items:
+        if "mypy" in item.nodeid:
+            # IMPORTANT: mypy type annotation tests leverage the pytest-mypy-plugins
+            # plugin, and are thus written in test_*.yml files.  As such, there are
+            # no explicit test functions on which we can apply a pytest.mark.mypy
+            # decorator.  Therefore, we mark them via this name-based, automatic
+            # marking approach, meaning that each test case must contain "mypy" in the
+            # name.
+            item.add_marker(pytest.mark.mypy)
 
 
 @pytest.fixture(autouse=True)
diff --git a/pyproject.toml b/pyproject.toml
@@ -44,8 +44,9 @@ dev = [
   "pytest",
   "pytest-cov",
   "pytest-env",
-  "pytest-xdist",
+  "pytest-mypy-plugins",
   "pytest-timeout",
+  "pytest-xdist",
   "ruff>=0.8.0",
   "sphinx",
   "sphinx_autosummary_accessors",
@@ -304,7 +305,12 @@ known-first-party = ["xarray"]
 ban-relative-imports = "all"
 
 [tool.pytest.ini_options]
-addopts = ["--strict-config", "--strict-markers"]
+addopts = [
+  "--strict-config",
+  "--strict-markers",
+  "--mypy-only-local-stub",
+  "--mypy-pyproject-toml-file=pyproject.toml",
+]
 
 # We want to forbid warnings from within xarray in our tests — instead we should
 # fix our own code, or mark the test itself as expecting a warning. So this:
@@ -361,6 +367,7 @@ filterwarnings = [
 log_cli_level = "INFO"
 markers = [
   "flaky: flaky tests",
+  "mypy: type annotation tests",
   "network: tests requiring a network connection",
   "slow: slow tests",
   "slow_hypothesis: slow hypothesis tests",
diff --git a/xarray/core/common.py b/xarray/core/common.py
@@ -6,7 +6,7 @@
 from contextlib import suppress
 from html import escape
 from textwrap import dedent
-from typing import TYPE_CHECKING, Any, TypeVar, Union, overload
+from typing import TYPE_CHECKING, Any, Concatenate, ParamSpec, TypeVar, Union, overload
 
 import numpy as np
 import pandas as pd
@@ -60,6 +60,7 @@
 T_Resample = TypeVar("T_Resample", bound="Resample")
 C = TypeVar("C")
 T = TypeVar("T")
+P = ParamSpec("P")
 
 
 class ImplementsArrayReduce:
@@ -718,11 +719,27 @@ def assign_attrs(self, *args: Any, **kwargs: Any) -> Self:
         out.attrs.update(*args, **kwargs)
         return out
 
+    @overload
+    def pipe(
+        self,
+        func: Callable[Concatenate[Self, P], T],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> T: ...
+
+    @overload
     def pipe(
         self,
-        func: Callable[..., T] | tuple[Callable[..., T], str],
+        func: tuple[Callable[..., T], str],
         *args: Any,
         **kwargs: Any,
+    ) -> T: ...
+
+    def pipe(
+        self,
+        func: Callable[Concatenate[Self, P], T] | tuple[Callable[P, T], str],
+        *args: P.args,
+        **kwargs: P.kwargs,
     ) -> T:
         """
         Apply ``func(self, *args, **kwargs)``
@@ -840,15 +857,19 @@ def pipe(
         pandas.DataFrame.pipe
         """
         if isinstance(func, tuple):
-            func, target = func
+            # Use different var when unpacking function from tuple because the type
+            # signature of the unpacked function differs from the expected type
+            # signature in the case where only a function is given, rather than a tuple.
+            # This makes type checkers happy at both call sites below.
+            f, target = func
             if target in kwargs:
                 raise ValueError(
                     f"{target} is both the pipe target and a keyword argument"
                 )
             kwargs[target] = self
-            return func(*args, **kwargs)
-        else:
-            return func(self, *args, **kwargs)
+            return f(*args, **kwargs)
+
+        return func(self, *args, **kwargs)
 
     def rolling_exp(
         self: T_DataWithCoords,
diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py
diff --git a/xarray/tests/test_dataarray_typing.yml b/xarray/tests/test_dataarray_typing.yml
diff --git a/xarray/tests/test_dataset_typing.yml b/xarray/tests/test_dataset_typing.yml
diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py
diff --git a/xarray/tests/test_datatree_typing.yml b/xarray/tests/test_datatree_typing.yml
