support passing a function to combine_attrs

doc/whats-new.rst

@@ -33,6 +33,8 @@ New Features
   By `Jimmy Westling <https://github.com/illviljan>`_.
 - Raise more informative error when decoding time variables with invalid reference dates.
   (:issue:`5199`, :pull:`5288`). By `Giacomo Caria <https://github.com/gcaria>`_.
+- allow passing a function to ``combine_attrs`` (:pull:`4896`).
+  By `Justus Magin <https://github.com/keewis>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/core/combine.py

@@ -432,8 +432,9 @@ def combine_nested(
           those of the first object with that dimension. Indexes for the same
           dimension must have the same size in all objects.
     combine_attrs : {"drop", "identical", "no_conflicts", "drop_conflicts", \
-                     "override"}, default: "drop"
-        String indicating how to combine attrs of the objects being merged:
+                     "override"} or callable, default: "drop"
+        A callable or a string indicating how to combine attrs of the objects being
+        merged:
 
         - "drop": empty attrs on returned Dataset.
         - "identical": all attrs must be the same on every object.
@@ -444,6 +445,10 @@ def combine_nested(
         - "override": skip comparing and copy attrs from the first dataset to
           the result.
 
+        If a callable, it must expect a sequence of ``attrs`` dicts as its only
+        parameter.
+
+
     Returns
     -------
     combined : xarray.Dataset
@@ -646,8 +651,9 @@ def combine_by_coords(
           those of the first object with that dimension. Indexes for the same
           dimension must have the same size in all objects.
     combine_attrs : {"drop", "identical", "no_conflicts", "drop_conflicts", \
-                     "override"}, default: "drop"
-        String indicating how to combine attrs of the objects being merged:
+                     "override"} or callable, default: "drop"
+        A callable or a string indicating how to combine attrs of the objects being
+        merged:
 
         - "drop": empty attrs on returned Dataset.
         - "identical": all attrs must be the same on every object.
@@ -658,6 +664,9 @@ def combine_by_coords(
         - "override": skip comparing and copy attrs from the first dataset to
           the result.
 
+        If a callable, it must expect a sequence of ``attrs`` dicts as its only
+        parameter.
+
     Returns
     -------
     combined : xarray.Dataset

xarray/core/concat.py

@@ -143,8 +143,9 @@ def concat(
           those of the first object with that dimension. Indexes for the same
           dimension must have the same size in all objects.
     combine_attrs : {"drop", "identical", "no_conflicts", "drop_conflicts", \
-                     "override"}, default: "override"
-        String indicating how to combine attrs of the objects being merged:
+                     "override"} or callable, default: "override"
+        A callable or a string indicating how to combine attrs of the objects being
+        merged:
 
         - "drop": empty attrs on returned Dataset.
         - "identical": all attrs must be the same on every object.
@@ -155,6 +156,9 @@ def concat(
         - "override": skip comparing and copy attrs from the first dataset to
           the result.
 
+        If a callable, it must expect a sequence of ``attrs`` dicts as its only
+        parameter.
+
     Returns
     -------
     concatenated : type of objs

xarray/core/merge.py

@@ -57,6 +57,11 @@
 )
 
 
+class Context:
+    def __init__(self, func):
+        self.func = func
+
+
 def broadcast_dimension_size(variables: List[Variable]) -> Dict[Hashable, int]:
     """Extract dimension sizes from a dictionary of variables.
 
@@ -502,13 +507,15 @@ def assert_valid_explicit_coords(variables, dims, explicit_coords):
             )
 
 
-def merge_attrs(variable_attrs, combine_attrs):
+def merge_attrs(variable_attrs, combine_attrs, context=None):
     """Combine attributes from different variables according to combine_attrs"""
     if not variable_attrs:
         # no attributes to merge
         return None
 
-    if combine_attrs == "drop":
+    if callable(combine_attrs):
+        return combine_attrs(variable_attrs, context=context)
+    elif combine_attrs == "drop":
         return {}
     elif combine_attrs == "override":
         return dict(variable_attrs[0])
@@ -585,7 +592,7 @@ def merge_core(
     join : {"outer", "inner", "left", "right"}, optional
         How to combine objects with different indexes.
     combine_attrs : {"drop", "identical", "no_conflicts", "drop_conflicts", \
-                     "override"}, optional
+                     "override"} or callable, default: "override"
         How to combine attributes of objects
     priority_arg : int, optional
         Optional argument in `objects` that takes precedence over the others.
@@ -696,8 +703,9 @@ def merge(
         variable names to fill values. Use a data array's name to
         refer to its values.
     combine_attrs : {"drop", "identical", "no_conflicts", "drop_conflicts", \
-                     "override"}, default: "override"
-        String indicating how to combine attrs of the objects being merged:
+                    "override"} or callable, default: "override"
+        A callable or a string indicating how to combine attrs of the objects being
+        merged:
 
         - "drop": empty attrs on returned Dataset.
         - "identical": all attrs must be the same on every object.
@@ -708,6 +716,9 @@ def merge(
         - "override": skip comparing and copy attrs from the first dataset to
           the result.
 
+        If a callable, it must expect a sequence of ``attrs`` dicts as its only
+        parameter.
+
     Returns
     -------
     Dataset

xarray/tests/test_combine.py

@@ -531,6 +531,9 @@ def test_auto_combine_2d_combine_attrs_kwarg(self):
         }
         expected_dict["override"] = expected.copy(deep=True)
         expected_dict["override"].attrs = {"a": 1}
+        f = lambda attrs, context: attrs[0]
+        expected_dict[f] = expected.copy(deep=True)
+        expected_dict[f].attrs = f([{"a": 1}], None)
 
         datasets = [[ds(0), ds(1), ds(2)], [ds(3), ds(4), ds(5)]]
 
@@ -714,6 +717,10 @@ def test_combine_coords_join_exact(self):
                 Dataset({"x": [0, 1], "y": [0, 1]}, attrs={"a": 1, "b": 2}),
             ),
             ("override", Dataset({"x": [0, 1], "y": [0, 1]}, attrs={"a": 1})),
+            (
+                lambda attrs, context: attrs[1],
+                Dataset({"x": [0, 1], "y": [0, 1]}, attrs={"a": 1, "b": 2}),
+            ),
         ],
     )
     def test_combine_coords_combine_attrs(self, combine_attrs, expected):

xarray/tests/test_concat.py

@@ -299,6 +299,13 @@ def test_concat_join_kwarg(self):
                 {"a": 41, "c": 43, "d": 44},
                 False,
             ),
+            (
+                lambda attrs, context: {"a": -1, "b": 0, "c": 1} if any(attrs) else {},
+                {"a": 41, "b": 42, "c": 43},
+                {"b": 2, "c": 43, "d": 44},
+                {"a": -1, "b": 0, "c": 1},
+                False,
+            ),
         ],
     )
     def test_concat_combine_attrs_kwarg(
@@ -354,6 +361,13 @@ def test_concat_combine_attrs_kwarg(
                 {"a": 41, "c": 43, "d": 44},
                 False,
             ),
+            (
+                lambda attrs, context: {"a": -1, "b": 0, "c": 1} if any(attrs) else {},
+                {"a": 41, "b": 42, "c": 43},
+                {"b": 2, "c": 43, "d": 44},
+                {"a": -1, "b": 0, "c": 1},
+                False,
+            ),
         ],
     )
     def test_concat_combine_attrs_kwarg_variables(

xarray/tests/test_merge.py

@@ -109,6 +109,13 @@ def test_merge_arrays_attrs_default(self):
                 {"a": 1, "c": np.array([3]), "d": 4},
                 False,
             ),
+            (
+                lambda attrs, context: attrs[1],
+                {"a": 1, "b": 2, "c": 3},
+                {"a": 4, "b": 3, "c": 1},
+                {"a": 4, "b": 3, "c": 1},
+                False,
+            ),
         ],
     )
     def test_merge_arrays_attrs(
@@ -161,6 +168,13 @@ def test_merge_arrays_attrs(
                 {"a": 1, "c": 3, "d": 4},
                 False,
             ),
+            (
+                lambda attrs, context: attrs[1],
+                {"a": 1, "b": 2, "c": 3},
+                {"a": 4, "b": 3, "c": 1},
+                {"a": 4, "b": 3, "c": 1},
+                False,
+            ),
         ],
     )
     def test_merge_arrays_attrs_variables(