Fix coordinate attr handling in xr.where(..., keep_attrs=True)

doc/whats-new.rst

@@ -34,7 +34,8 @@ Deprecations
 
 Bug fixes
 ~~~~~~~~~
-
+- Fix handling of coordinate attributes in :py:func:`where`. (:issue:`7220`, :pull:`7229`)
+  By `Sam Levang <https://github.com/slevang>`_.
 - Import ``nc_time_axis`` when needed (:issue:`7275`, :pull:`7276`).
   By `Michael Niklas <https://github.com/headtr1ck>`_.
 

xarray/core/computation.py

@@ -25,7 +25,7 @@
 import numpy as np
 
 from . import dtypes, duck_array_ops, utils
-from .alignment import align, deep_align
+from .alignment import align, broadcast, deep_align
 from .common import zeros_like
 from .duck_array_ops import datetime_to_numeric
 from .indexes import Index, filter_indexes_from_coords
@@ -1855,12 +1855,19 @@ def where(cond, x, y, keep_attrs=None):
     Dataset.where, DataArray.where :
         equivalent methods
     """
+    from .dataarray import DataArray
+
     if keep_attrs is None:
         keep_attrs = _get_keep_attrs(default=False)
     if keep_attrs is True:
         # keep the attributes of x, the second parameter, by default to
         # be consistent with the `where` method of `DataArray` and `Dataset`
-        keep_attrs = lambda attrs, context: getattr(x, "attrs", {})
+        keep_attrs = lambda attrs, context: attrs[1]
+        # cast non-xarray objects to DataArray to get empty attrs
+        cond, x, y = (v if hasattr(v, "attrs") else DataArray(v) for v in [cond, x, y])
+        # explicitly broadcast to ensure we also get empty coord attrs
+        # take coord attrs preferentially from x, then y, then cond
+        x, y, cond = broadcast(x, y, cond)
 
     # alignment for three arguments is complicated, so don't support it yet
     return apply_ufunc(

xarray/tests/test_computation.py

@@ -1925,16 +1925,42 @@ def test_where() -> None:
 
 
 def test_where_attrs() -> None:
-    cond = xr.DataArray([True, False], dims="x", attrs={"attr": "cond"})
-    x = xr.DataArray([1, 1], dims="x", attrs={"attr": "x"})
-    y = xr.DataArray([0, 0], dims="x", attrs={"attr": "y"})
+    cond = xr.DataArray([True, False], coords={"x": [0, 1]}, attrs={"attr": "cond_da"})
+    cond["x"].attrs = {"attr": "cond_coord"}
+    x = xr.DataArray([1, 1], coords={"x": [0, 1]}, attrs={"attr": "x_da"})
+    x["x"].attrs = {"attr": "x_coord"}
+    y = xr.DataArray([0, 0], coords={"x": [0, 1]}, attrs={"attr": "y_da"})
+    y["x"].attrs = {"attr": "y_coord"}
+
+    # 3 DataArrays, takes attrs from x
     actual = xr.where(cond, x, y, keep_attrs=True)
-    expected = xr.DataArray([1, 0], dims="x", attrs={"attr": "x"})
+    expected = xr.DataArray([1, 0], coords={"x": [0, 1]}, attrs={"attr": "x_da"})
+    expected["x"].attrs = {"attr": "x_coord"}
     assert_identical(expected, actual)
 
-    # ensure keep_attrs can handle scalar values
+    # x as a scalar, takes coord attrs only from y
+    actual = xr.where(cond, 0, y, keep_attrs=True)
+    expected = xr.DataArray([0, 0], coords={"x": [0, 1]})
+    expected["x"].attrs = {"attr": "y_coord"}
+    assert_identical(expected, actual)
+
+    # y as a scalar, takes attrs from x
+    actual = xr.where(cond, x, 0, keep_attrs=True)
+    expected = xr.DataArray([1, 0], coords={"x": [0, 1]}, attrs={"attr": "x_da"})
+    expected["x"].attrs = {"attr": "x_coord"}
+    assert_identical(expected, actual)
+
+    # x and y as a scalar, takes coord attrs only from cond
     actual = xr.where(cond, 1, 0, keep_attrs=True)
-    assert actual.attrs == {}
+    expected = xr.DataArray([1, 0], coords={"x": [0, 1]})
+    expected["x"].attrs = {"attr": "cond_coord"}
+    assert_identical(expected, actual)
+
+    # cond and y as a scalar, takes attrs from x
+    actual = xr.where(True, x, y, keep_attrs=True)
+    expected = xr.DataArray([1, 1], coords={"x": [0, 1]}, attrs={"attr": "x_da"})
+    expected["x"].attrs = {"attr": "x_coord"}
+    assert_identical(expected, actual)
 
 
 @pytest.mark.parametrize(