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

doc/whats-new.rst

@@ -61,6 +61,8 @@ Bug fixes
   now reopens the file from scratch for h5netcdf and scipy netCDF backends,
   rather than reusing a cached version (:issue:`4240`, :issue:`4862`).
   By `Stephan Hoyer <https://github.com/shoyer>`_.
+- Fix handling of coordinate attributes in ``xarray.where``. (:issue:`7220`, :pull:`7229`)
+  By `Sam Levang <https://github.com/slevang>`_.
 - Fixed bug where :py:meth:`Dataset.coarsen.construct` would demote non-dimension coordinates to variables. (:pull:`7233`)
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
 - Raise a TypeError when trying to plot empty data (:issue:`7156`, :pull:`7228`).

xarray/core/computation.py

@@ -1789,7 +1789,8 @@ def where(cond, x, y, keep_attrs=None):
     y : scalar, array, Variable, DataArray or Dataset
         values to choose from where `cond` is False
     keep_attrs : bool or str or callable, optional
-        How to treat attrs. If True, keep the attrs of `x`.
+        How to treat attrs. If True, keep the attrs of `x`,
+        unless `x` is a scalar, then keep the attrs of `y`.
 
     Returns
     -------
@@ -1860,7 +1861,7 @@ def where(cond, x, y, keep_attrs=None):
     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] if len(attrs) > 1 else {}
 
     # alignment for three arguments is complicated, so don't support it yet
     return apply_ufunc(

xarray/tests/test_computation.py

@@ -1925,16 +1925,36 @@ 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 attrs from y
+    actual = xr.where(cond, 0, y, keep_attrs=True)
+    expected = xr.DataArray([0, 0], coords={"x": [0, 1]}, attrs={"attr": "y_da"})
+    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)
 
 
 @pytest.mark.parametrize(