More precisely type pipe methods.

In addition, enhance mypy job configuration to support running it
locally via `act`.

Fixes #9997

Author: chuckwondo

.github/workflows/ci-additional.yaml

@@ -33,6 +33,15 @@ jobs:
         with:
           keyword: "[skip-ci]"
 
+  detect-act:
+    name: Detect 'act' runner
+    runs-on: ubuntu-latest
+    outputs:
+      running: ${{ steps.detect-act.outputs.running }}
+    steps:
+      - id: detect-act
+        run: echo "running=${{ env.ACT || 'false' }}" >> "$GITHUB_OUTPUT"
+
   doctest:
     name: Doctests
     runs-on: "ubuntu-latest"
@@ -81,15 +90,23 @@ jobs:
           python -m pytest --doctest-modules xarray --ignore xarray/tests -Werror
 
   mypy:
-    name: Mypy
+    strategy:
+      matrix:
+        include:
+          - python-version: "3.10"
+            codecov-flags: mypy-min
+          - python-version: "3.12"
+            codecov-flags: mypy
+    name: Mypy ${{ matrix.python-version }}
     runs-on: "ubuntu-latest"
-    needs: detect-ci-trigger
+    needs: [detect-ci-trigger, detect-act]
+    if: always() && (needs.detect-ci-trigger.outputs.triggered == 'true' || needs.detect-act.outputs.running == 'true')
     defaults:
       run:
         shell: bash -l {0}
     env:
       CONDA_ENV_FILE: ci/requirements/environment.yml
-      PYTHON_VERSION: "3.12"
+      PYTHON_VERSION: ${{ matrix.python-version }}
 
     steps:
       - uses: actions/checkout@v4
@@ -116,68 +133,23 @@ jobs:
           python xarray/util/print_versions.py
       - name: Install mypy
         run: |
-          python -m pip install mypy
+          python -m pip install mypy pytest-mypy-plugins
 
       - name: Run mypy
         run: |
           python -m mypy --install-types --non-interactive --cobertura-xml-report mypy_report
 
-      - name: Upload mypy coverage to Codecov
-        uses: codecov/[email protected]
-        with:
-          files: mypy_report/cobertura.xml
-          flags: mypy
-          env_vars: PYTHON_VERSION
-          name: codecov-umbrella
-          fail_ci_if_error: false
-
-  mypy-min:
-    name: Mypy 3.10
-    runs-on: "ubuntu-latest"
-    needs: detect-ci-trigger
-    defaults:
-      run:
-        shell: bash -l {0}
-    env:
-      CONDA_ENV_FILE: ci/requirements/environment.yml
-      PYTHON_VERSION: "3.10"
-
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          fetch-depth: 0 # Fetch all history for all branches and tags.
-
-      - name: set environment variables
+      - name: Run mypy tests
+        # Run pytest with mypy plugin even if mypy analysis in previous step fails.
+        if: ${{ always() }}
         run: |
-          echo "TODAY=$(date +'%Y-%m-%d')" >> $GITHUB_ENV
-      - name: Setup micromamba
-        uses: mamba-org/setup-micromamba@v2
-        with:
-          environment-file: ${{env.CONDA_ENV_FILE}}
-          environment-name: xarray-tests
-          create-args: >-
-            python=${{env.PYTHON_VERSION}}
-          cache-environment: true
-          cache-environment-key: "${{runner.os}}-${{runner.arch}}-py${{env.PYTHON_VERSION}}-${{env.TODAY}}-${{hashFiles(env.CONDA_ENV_FILE)}}"
-      - name: Install xarray
-        run: |
-          python -m pip install --no-deps -e .
-      - name: Version info
-        run: |
-          python xarray/util/print_versions.py
-      - name: Install mypy
-        run: |
-          python -m pip install mypy
-
-      - name: Run mypy
-        run: |
-          python -m mypy --install-types --non-interactive --cobertura-xml-report mypy_report
+          python -m pytest -v --mypy-only-local-stub --mypy-pyproject-toml-file=pyproject.toml xarray/**/test_*.yml
 
       - name: Upload mypy coverage to Codecov
         uses: codecov/[email protected]
         with:
           files: mypy_report/cobertura.xml
-          flags: mypy-min
+          flags: ${{ matrix.codecov-flags }}
           env_vars: PYTHON_VERSION
           name: codecov-umbrella
           fail_ci_if_error: false

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,

xarray/core/datatree.py

@@ -12,7 +12,17 @@
     Mapping,
 )
 from html import escape
-from typing import TYPE_CHECKING, Any, Literal, NoReturn, Union, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Concatenate,
+    Literal,
+    NoReturn,
+    ParamSpec,
+    TypeVar,
+    Union,
+    overload,
+)
 
 from xarray.core import utils
 from xarray.core._aggregations import DataTreeAggregations
@@ -79,18 +89,23 @@
 # """
 # DEVELOPERS' NOTE
 # ----------------
-# The idea of this module is to create a `DataTree` class which inherits the tree structure from TreeNode, and also copies
-# the entire API of `xarray.Dataset`, but with certain methods decorated to instead map the dataset function over every
-# node in the tree. As this API is copied without directly subclassing `xarray.Dataset` we instead create various Mixin
-# classes (in ops.py) which each define part of `xarray.Dataset`'s extensive API.
+# The idea of this module is to create a `DataTree` class which inherits the tree
+# structure from TreeNode, and also copies the entire API of `xarray.Dataset`, but with
+# certain methods decorated to instead map the dataset function over every node in the
+# tree. As this API is copied without directly subclassing `xarray.Dataset` we instead
+# create various Mixin classes (in ops.py) which each define part of `xarray.Dataset`'s
+# extensive API.
 #
-# Some of these methods must be wrapped to map over all nodes in the subtree. Others are fine to inherit unaltered
-# (normally because they (a) only call dataset properties and (b) don't return a dataset that should be nested into a new
-# tree) and some will get overridden by the class definition of DataTree.
+# Some of these methods must be wrapped to map over all nodes in the subtree. Others are
+# fine to inherit unaltered (normally because they (a) only call dataset properties and
+# (b) don't return a dataset that should be nested into a new tree) and some will get
+# overridden by the class definition of DataTree.
 # """
 
 
 T_Path = Union[str, NodePath]
+T = TypeVar("T")
+P = ParamSpec("P")
 
 
 def _collect_data_and_coord_variables(
@@ -1460,9 +1475,28 @@ def map_over_datasets(
         # TODO fix this typing error
         return map_over_datasets(func, self, *args)
 
+    @overload
+    def pipe(
+        self,
+        func: Callable[Concatenate[Self, P], T],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> T: ...
+
+    @overload
+    def pipe(
+        self,
+        func: tuple[Callable[..., T], str],
+        *args: Any,
+        **kwargs: Any,
+    ) -> T: ...
+
     def pipe(
-        self, func: Callable | tuple[Callable, str], *args: Any, **kwargs: Any
-    ) -> Any:
+        self,
+        func: Callable[Concatenate[Self, P], T] | tuple[Callable[..., T], str],
+        *args: Any,
+        **kwargs: Any,
+    ) -> T:
         """Apply ``func(self, *args, **kwargs)``
 
         This method replicates the pandas method of the same name.
@@ -1482,7 +1516,7 @@ def pipe(
 
         Returns
         -------
-        object : Any
+        object : T
             the return type of ``func``.
 
         Notes
@@ -1510,15 +1544,19 @@ def 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
-        else:
-            args = (self,) + args
-        return func(*args, **kwargs)
+            return f(*args, **kwargs)
+
+        return func(self, *args, **kwargs)
 
     # TODO some kind of .collapse() or .flatten() method to merge a subtree