add kwargs to map_over_datasets (similar to apply_ufunc), add test.

kmuehlbauer · kmuehlbauer · commit 3c07cceee232 · 2025-01-31T11:52:27.000+01:00
diff --git a/xarray/core/datatree_mapping.py b/xarray/core/datatree_mapping.py
@@ -2,6 +2,7 @@
 
 import sys
 from collections.abc import Callable, Mapping
+from functools import partial
 from typing import TYPE_CHECKING, Any, cast, overload
 
 from xarray.core.dataset import Dataset
@@ -31,7 +32,9 @@ def map_over_datasets(
 
 
 def map_over_datasets(
-    func: Callable[..., Dataset | None | tuple[Dataset | None, ...]], *args: Any
+    func: Callable[..., Dataset | None | tuple[Dataset | None, ...]],
+    *args: Any,
+    kwargs: Mapping | None = None,
 ) -> DataTree | tuple[DataTree, ...]:
     """
     Applies a function to every dataset in one or more DataTree objects with
@@ -68,6 +71,8 @@ def map_over_datasets(
     *args : tuple, optional
         Positional arguments passed on to `func`. Any DataTree arguments will be
         converted to Dataset objects via `.dataset`.
+    kwargs : dict, optional
+        Optional keyword arguments passed directly on to call ``func``.
 
     Returns
     -------
@@ -85,6 +90,12 @@ def map_over_datasets(
 
     from xarray.core.datatree import DataTree
 
+    if kwargs is None:
+        kwargs = {}
+
+    if kwargs:
+        func = partial(func, **kwargs)
+
     # Walk all trees simultaneously, applying func to all nodes that lie in same position in different trees
     # We don't know which arguments are DataTrees so we zip all arguments together as iterables
     # Store tuples of results in a dict because we don't yet know how many trees we need to rebuild to return
diff --git a/xarray/tests/test_datatree_mapping.py b/xarray/tests/test_datatree_mapping.py
@@ -51,6 +51,19 @@ def test_single_tree_arg_plus_arg(self, create_test_datatree):
         result_tree = map_over_datasets(lambda x, y: x * y, 10.0, dt)
         assert_equal(result_tree, expected)
 
+    def test_single_tree_arg_plus_kwarg(self, create_test_datatree):
+        dt = create_test_datatree()
+        expected = create_test_datatree(modify=lambda ds: (10.0 * ds))
+
+        def multiply_by_kwarg(ds, **kwargs):
+            ds = ds * kwargs.pop("multiplier")
+            return ds
+
+        result_tree = map_over_datasets(
+            multiply_by_kwarg, dt, kwargs=dict(multiplier=10.0)
+        )
+        assert_equal(result_tree, expected)
+
     def test_multiple_tree_args(self, create_test_datatree):
         dt1 = create_test_datatree()
         dt2 = create_test_datatree()
