From 023ce94d3c22999cd5de4a3900a2067e387e603e Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Wed, 11 Dec 2019 22:42:01 +0100
Subject: [PATCH 01/19] get the align tests to pass

---
 xarray/core/variable.py    |  2 +-
 xarray/tests/test_units.py | 72 +++++++++++++++++++++++---------------
 2 files changed, 44 insertions(+), 30 deletions(-)

diff --git a/xarray/core/variable.py b/xarray/core/variable.py
index aa04cffb5ea..cddd4c8a6e1 100644
--- a/xarray/core/variable.py
+++ b/xarray/core/variable.py
@@ -739,7 +739,7 @@ def _getitem_with_mask(self, key, fill_value=dtypes.NA):
 
             data = as_indexable(self._data)[actual_indexer]
             mask = indexing.create_mask(indexer, self.shape, data)
-            data = duck_array_ops.where(mask, fill_value, data)
+            data = duck_array_ops.where(~mask, data, fill_value)
         else:
             # array cannot be indexed along dimensions of size 0, so just
             # build the mask directly instead.
diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index f8a8a259c1f..563fb3bb59a 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -380,7 +380,6 @@ def test_apply_ufunc_dataset(dtype):
     assert_equal_with_units(expected, actual)
 
 
-@pytest.mark.xfail(reason="blocked by `reindex` / `where`")
 @pytest.mark.parametrize(
     "unit,error",
     (
@@ -402,36 +401,40 @@ def test_apply_ufunc_dataset(dtype):
         "coords",
     ),
 )
-@pytest.mark.parametrize("fill_value", (np.float64(10), np.float64(np.nan)))
+@pytest.mark.parametrize("fill_value", (10, np.nan))
 def test_align_dataarray(fill_value, variant, unit, error, dtype):
     original_unit = unit_registry.m
 
     variants = {
-        "data": (unit, 1, 1),
-        "dims": (original_unit, unit, 1),
-        "coords": (original_unit, 1, unit),
+        "data": (unit, original_unit, original_unit),
+        "dims": (original_unit, unit, original_unit),
+        "coords": (original_unit, original_unit, unit),
     }
     data_unit, dim_unit, coord_unit = variants.get(variant)
 
     array1 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * original_unit
     array2 = np.linspace(0, 8, 2 * 5).reshape(2, 5).astype(dtype) * data_unit
     x = np.arange(2) * original_unit
-    x_a1 = np.array([10, 5]) * original_unit
-    x_a2 = np.array([10, 5]) * coord_unit
 
     y1 = np.arange(5) * original_unit
     y2 = np.arange(2, 7) * dim_unit
+    y_a1 = np.array([3, 5, 7, 8, 9]) * original_unit
+    y_a2 = np.array([7, 8, 9, 11, 13]) * coord_unit
 
-    data_array1 = xr.DataArray(
-        data=array1, coords={"x": x, "x_a": ("x", x_a1), "y": y1}, dims=("x", "y")
-    )
-    data_array2 = xr.DataArray(
-        data=array2, coords={"x": x, "x_a": ("x", x_a2), "y": y2}, dims=("x", "y")
-    )
+    coords1 = {"x": x, "y": y1}
+    coords2 = {"x": x, "y": y2}
+    if variant == "coords":
+        coords1["y_a"] = ("y", y_a1)
+        coords2["y_a"] = ("y", y_a2)
+
+    data_array1 = xr.DataArray(data=array1, coords=coords1, dims=("x", "y"))
+    data_array2 = xr.DataArray(data=array2, coords=coords2, dims=("x", "y"))
 
     fill_value = fill_value * data_unit
     func = function(xr.align, join="outer", fill_value=fill_value)
-    if error is not None:
+    if error is not None and not (
+        np.isnan(fill_value) and not isinstance(fill_value, Quantity)
+    ):
         with pytest.raises(error):
             func(data_array1, data_array2)
 
@@ -451,7 +454,10 @@ def test_align_dataarray(fill_value, variant, unit, error, dtype):
         **stripped_kwargs,
     )
     expected_a = attach_units(expected_a, units_a)
-    expected_b = convert_units(attach_units(expected_b, units_a), units_b)
+    if isinstance(array2, Quantity):
+        expected_b = convert_units(attach_units(expected_b, units_a), units_b)
+    else:
+        expected_b = attach_units(expected_b, units_b)
 
     actual_a, actual_b = func(data_array1, data_array2)
 
@@ -459,7 +465,6 @@ def test_align_dataarray(fill_value, variant, unit, error, dtype):
     assert_equal_with_units(expected_b, actual_b)
 
 
-@pytest.mark.xfail(reason="blocked by `reindex` / `where`")
 @pytest.mark.parametrize(
     "unit,error",
     (
@@ -485,31 +490,37 @@ def test_align_dataarray(fill_value, variant, unit, error, dtype):
 def test_align_dataset(fill_value, unit, variant, error, dtype):
     original_unit = unit_registry.m
 
-    variants = {"data": (unit, 1, 1), "dims": (1, unit, 1), "coords": (1, 1, unit)}
+    variants = {
+        "data": (unit, original_unit, original_unit),
+        "dims": (original_unit, unit, original_unit),
+        "coords": (original_unit, original_unit, unit),
+    }
     data_unit, dim_unit, coord_unit = variants.get(variant)
 
     array1 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * original_unit
     array2 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * data_unit
 
     x = np.arange(2) * original_unit
-    x_a1 = np.array([10, 5]) * original_unit
-    x_a2 = np.array([10, 5]) * coord_unit
 
     y1 = np.arange(5) * original_unit
     y2 = np.arange(2, 7) * dim_unit
+    y_a1 = np.array([3, 5, 7, 8, 9]) * original_unit
+    y_a2 = np.array([7, 8, 9, 11, 13]) * coord_unit
 
-    ds1 = xr.Dataset(
-        data_vars={"a": (("x", "y"), array1)},
-        coords={"x": x, "x_a": ("x", x_a1), "y": y1},
-    )
-    ds2 = xr.Dataset(
-        data_vars={"a": (("x", "y"), array2)},
-        coords={"x": x, "x_a": ("x", x_a2), "y": y2},
-    )
+    coords1 = {"x": x, "y": y1}
+    coords2 = {"x": x, "y": y2}
+    if variant == "coords":
+        coords1["y_a"] = ("y", y_a1)
+        coords2["y_a"] = ("y", y_a2)
+
+    ds1 = xr.Dataset(data_vars={"a": (("x", "y"), array1)}, coords=coords1)
+    ds2 = xr.Dataset(data_vars={"a": (("x", "y"), array2)}, coords=coords2)
 
     fill_value = fill_value * data_unit
     func = function(xr.align, join="outer", fill_value=fill_value)
-    if error is not None:
+    if error is not None and not (
+        np.isnan(fill_value) and not isinstance(fill_value, Quantity)
+    ):
         with pytest.raises(error):
             func(ds1, ds2)
 
@@ -527,7 +538,10 @@ def test_align_dataset(fill_value, unit, variant, error, dtype):
         strip_units(ds1), strip_units(convert_units(ds2, units_a)), **stripped_kwargs
     )
     expected_a = attach_units(expected_a, units_a)
-    expected_b = convert_units(attach_units(expected_b, units_a), units_b)
+    if isinstance(array2, Quantity):
+        expected_b = convert_units(attach_units(expected_b, units_a), units_b)
+    else:
+        expected_b = attach_units(expected_b, units_b)
 
     actual_a, actual_b = func(ds1, ds2)
 

From 9dba9e367fddd25a927dbcc7d5ba405c9322e702 Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Wed, 11 Dec 2019 22:45:16 +0100
Subject: [PATCH 02/19] add pint to the upstream-dev ci job

---
 ci/azure/install.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/ci/azure/install.yml b/ci/azure/install.yml
index e4f3a0b9e16..0c10d82e9fc 100644
--- a/ci/azure/install.yml
+++ b/ci/azure/install.yml
@@ -27,6 +27,7 @@ steps:
         git+https://github.com/zarr-developers/zarr \
         git+https://github.com/Unidata/cftime \
         git+https://github.com/mapbox/rasterio \
+        git+https://github.com/hgrecco/pint \
         git+https://github.com/pydata/bottleneck
   condition: eq(variables['UPSTREAM_DEV'], 'true')
   displayName: Install upstream dev dependencies

From ffa9de3e4ce3ec62a087bb109f75bb0e55423a88 Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Thu, 12 Dec 2019 09:38:01 +0100
Subject: [PATCH 03/19] special case for booleans

---
 xarray/core/variable.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/xarray/core/variable.py b/xarray/core/variable.py
index cddd4c8a6e1..2a1d804e727 100644
--- a/xarray/core/variable.py
+++ b/xarray/core/variable.py
@@ -739,7 +739,11 @@ def _getitem_with_mask(self, key, fill_value=dtypes.NA):
 
             data = as_indexable(self._data)[actual_indexer]
             mask = indexing.create_mask(indexer, self.shape, data)
-            data = duck_array_ops.where(~mask, data, fill_value)
+            if isinstance(mask, bool):
+                mask = not mask
+            else:
+                mask = ~mask
+            data = duck_array_ops.where(mask, data, fill_value)
         else:
             # array cannot be indexed along dimensions of size 0, so just
             # build the mask directly instead.

From 9aa07974fe5d7263fdf62a95c4778560d235ff3b Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Thu, 12 Dec 2019 11:25:49 +0100
Subject: [PATCH 04/19] silence the pint behaviour change warning

---
 xarray/tests/test_units.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 563fb3bb59a..5d2447e5e16 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -1,4 +1,5 @@
 import operator
+import warnings
 
 import numpy as np
 import pandas as pd
@@ -11,6 +12,9 @@
 pint = pytest.importorskip("pint")
 DimensionalityError = pint.errors.DimensionalityError
 
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+    pint.Quantity([])
 
 unit_registry = pint.UnitRegistry()
 Quantity = unit_registry.Quantity

From bb93ea6ab23be0637e49351c6e8430c7189eb1bf Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Thu, 12 Dec 2019 11:27:11 +0100
Subject: [PATCH 05/19] preprocess the unit mapping parameter to convert_units

---
 xarray/tests/test_units.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 5d2447e5e16..a4125bd7b1c 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -207,6 +207,11 @@ def attach_units(obj, units):
 
 
 def convert_units(obj, to):
+    # preprocess
+    to = {
+        key: None if not isinstance(value, unit_registry.Unit) else value
+        for key, value in to.items()
+    }
     if isinstance(obj, xr.Dataset):
         data_vars = {
             name: convert_units(array.variable, {None: to.get(name)})

From 5ce66cdd3044b6e02c27525ec3adbd286cd76812 Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Thu, 12 Dec 2019 11:29:48 +0100
Subject: [PATCH 06/19] use assert_allclose and assert_identical instead

---
 xarray/tests/test_units.py | 39 ++++++++++++++++++++++----------------
 1 file changed, 23 insertions(+), 16 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index a4125bd7b1c..a3873d7bf94 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -8,6 +8,7 @@
 import xarray as xr
 from xarray.core import formatting
 from xarray.core.npcompat import IS_NEP18_ACTIVE
+from xarray.testing import assert_allclose, assert_identical
 
 pint = pytest.importorskip("pint")
 DimensionalityError = pint.errors.DimensionalityError
@@ -364,7 +365,7 @@ def test_apply_ufunc_dataarray(dtype):
     expected = attach_units(func(strip_units(data_array)), extract_units(data_array))
     actual = func(data_array)
 
-    assert_equal_with_units(expected, actual)
+    assert_identical(expected, actual)
 
 
 def test_apply_ufunc_dataset(dtype):
@@ -386,7 +387,7 @@ def test_apply_ufunc_dataset(dtype):
     expected = attach_units(func(strip_units(ds)), extract_units(ds))
     actual = func(ds)
 
-    assert_equal_with_units(expected, actual)
+    assert_identical(expected, actual)
 
 
 @pytest.mark.parametrize(
@@ -470,8 +471,10 @@ def test_align_dataarray(fill_value, variant, unit, error, dtype):
 
     actual_a, actual_b = func(data_array1, data_array2)
 
-    assert_equal_with_units(expected_a, actual_a)
-    assert_equal_with_units(expected_b, actual_b)
+    assert_allclose(expected_a, actual_a)
+    assert extract_units(expected_a) == extract_units(actual_a)
+    assert_allclose(expected_b, actual_b)
+    assert extract_units(expected_b) == extract_units(actual_b)
 
 
 @pytest.mark.parametrize(
@@ -554,8 +557,10 @@ def test_align_dataset(fill_value, unit, variant, error, dtype):
 
     actual_a, actual_b = func(ds1, ds2)
 
-    assert_equal_with_units(expected_a, actual_a)
-    assert_equal_with_units(expected_b, actual_b)
+    assert_allclose(expected_a, actual_a)
+    assert extract_units(expected_a) == extract_units(actual_a)
+    assert_allclose(expected_b, actual_b)
+    assert extract_units(expected_b) == extract_units(actual_b)
 
 
 def test_broadcast_dataarray(dtype):
@@ -571,8 +576,8 @@ def test_broadcast_dataarray(dtype):
     )
     actual_a, actual_b = xr.broadcast(a, b)
 
-    assert_equal_with_units(expected_a, actual_a)
-    assert_equal_with_units(expected_b, actual_b)
+    assert_identical(expected_a, actual_a)
+    assert_identical(expected_b, actual_b)
 
 
 def test_broadcast_dataset(dtype):
@@ -656,7 +661,7 @@ def test_combine_by_coords(variant, unit, error, dtype):
     )
     actual = xr.combine_by_coords([ds, other])
 
-    assert_equal_with_units(expected, actual)
+    assert_identical(expected, actual)
 
 
 @pytest.mark.parametrize(
@@ -760,7 +765,7 @@ def test_combine_nested(variant, unit, error, dtype):
     )
     actual = func([[ds1, ds2], [ds3, ds4]])
 
-    assert_equal_with_units(expected, actual)
+    assert_identical(expected, actual)
 
 
 @pytest.mark.parametrize(
@@ -812,7 +817,7 @@ def test_concat_dataarray(variant, unit, error, dtype):
     )
     actual = xr.concat([arr1, arr2], dim="x")
 
-    assert_equal_with_units(expected, actual)
+    assert_identical(expected, actual)
 
 
 @pytest.mark.parametrize(
@@ -862,7 +867,7 @@ def test_concat_dataset(variant, unit, error, dtype):
     )
     actual = xr.concat([ds1, ds2], dim="x")
 
-    assert_equal_with_units(expected, actual)
+    assert_identical(expected, actual)
 
 
 @pytest.mark.xfail(reason="blocked by `reindex` / `where`")
@@ -952,7 +957,8 @@ def test_merge_dataarray(variant, unit, error, dtype):
     )
     actual = func([arr1, arr2, arr3])
 
-    assert_equal_with_units(expected, actual)
+    assert extract_units(expected) == extract_units(actual)
+    assert_allclose(expected, actual)
 
 
 @pytest.mark.xfail(reason="blocked by `reindex` / `where`")
@@ -1035,7 +1041,8 @@ def test_merge_dataset(variant, unit, error, dtype):
     )
     actual = func([ds1, ds2, ds3])
 
-    assert_equal_with_units(expected, actual)
+    assert extract_units(expected) == extract_units(actual)
+    assert_allclose(expected, actual)
 
 
 @pytest.mark.parametrize("func", (xr.zeros_like, xr.ones_like))
@@ -1047,7 +1054,7 @@ def test_replication_dataarray(func, dtype):
     expected = xr.DataArray(data=numpy_func(array), dims="x")
     actual = func(data_array)
 
-    assert_equal_with_units(expected, actual)
+    assert_identical(expected, actual)
 
 
 @pytest.mark.parametrize("func", (xr.zeros_like, xr.ones_like))
@@ -1069,7 +1076,7 @@ def test_replication_dataset(func, dtype):
     )
     actual = func(ds)
 
-    assert_equal_with_units(expected, actual)
+    assert_identical(expected, actual)
 
 
 @pytest.mark.xfail(

From cad1308b980c539afb335cc1a63b134302e0e5df Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Thu, 12 Dec 2019 11:31:24 +0100
Subject: [PATCH 07/19] clean up a few tests

---
 xarray/tests/test_units.py | 98 ++++++++++++++++++++++++++++----------
 1 file changed, 73 insertions(+), 25 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index a3873d7bf94..cb01ec049e4 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -570,10 +570,12 @@ def test_broadcast_dataarray(dtype):
     a = xr.DataArray(data=array1, dims="x")
     b = xr.DataArray(data=array2, dims="y")
 
-    expected_a, expected_b = tuple(
-        attach_units(elem, extract_units(a))
-        for elem in xr.broadcast(strip_units(a), strip_units(b))
-    )
+    units_a = extract_units(a)
+    units_b = extract_units(b)
+    expected_a, expected_b = xr.broadcast(strip_units(a), strip_units(b))
+    expected_a = attach_units(expected_a, units_a)
+    expected_b = convert_units(attach_units(expected_b, units_a), units_b)
+
     actual_a, actual_b = xr.broadcast(a, b)
 
     assert_identical(expected_a, actual_a)
@@ -584,14 +586,33 @@ def test_broadcast_dataset(dtype):
     array1 = np.linspace(0, 10, 2) * unit_registry.Pa
     array2 = np.linspace(0, 10, 3) * unit_registry.Pa
 
-    ds = xr.Dataset(data_vars={"a": ("x", array1), "b": ("y", array2)})
+    x1 = np.arange(2)
+    y1 = np.arange(3)
+
+    x2 = np.arange(2, 4)
+    y2 = np.arange(3, 6)
 
-    (expected,) = tuple(
-        attach_units(elem, extract_units(ds)) for elem in xr.broadcast(strip_units(ds))
+    ds = xr.Dataset(
+        data_vars={"a": ("x", array1), "b": ("y", array2)}, coords={"x": x1, "y": y1}
+    )
+    other = xr.Dataset(
+        data_vars={
+            "a": ("x", array1.to(unit_registry.hPa)),
+            "b": ("y", array2.to(unit_registry.hPa)),
+        },
+        coords={"x": x2, "y": y2},
     )
-    (actual,) = xr.broadcast(ds)
 
-    assert_equal_with_units(expected, actual)
+    units_a = extract_units(ds)
+    units_b = extract_units(other)
+    expected_a, expected_b = xr.broadcast(strip_units(ds), strip_units(other))
+    expected_a = attach_units(expected_a, units_a)
+    expected_b = attach_units(expected_b, units_b)
+
+    actual_a, actual_b = xr.broadcast(ds, other)
+
+    assert_identical(expected_a, actual_a)
+    assert_identical(expected_b, actual_b)
 
 
 @pytest.mark.parametrize(
@@ -943,18 +964,33 @@ def test_merge_dataarray(variant, unit, error, dtype):
     )
 
     func = function(xr.merge)
-    if error is not None:
+    if error is not None and variant != "data":
         with pytest.raises(error):
             func([arr1, arr2, arr3])
 
         return
 
-    units = {name: original_unit for name in list("abcuvwxyz")}
+    units = {name: original_unit for name in list("uvwxyz")}
     convert_and_strip = lambda arr: strip_units(convert_units(arr, units))
-    expected = attach_units(
-        func([strip_units(arr1), convert_and_strip(arr2), convert_and_strip(arr3)]),
-        units,
+    expected_units = {
+        "a": original_unit,
+        "b": data_unit,
+        "c": data_unit,
+        "u": coord_unit,
+        "v": coord_unit,
+        "w": coord_unit,
+        "x": original_unit,
+        "y": original_unit,
+        "z": dim_unit,
+    }
+    expected = convert_units(
+        attach_units(
+            func([strip_units(arr1), convert_and_strip(arr2), convert_and_strip(arr3)]),
+            {**units, **{"a": original_unit, "b": data_unit, "c": data_unit}},
+        ),
+        expected_units,
     )
+
     actual = func([arr1, arr2, arr3])
 
     assert extract_units(expected) == extract_units(actual)
@@ -1002,7 +1038,7 @@ def test_merge_dataset(variant, unit, error, dtype):
 
     ds1 = xr.Dataset(
         data_vars={"a": (("y", "x"), array1), "b": (("y", "x"), array2)},
-        coords={"x": x, "y": y, "z": ("x", z)},
+        coords={"x": x, "y": y, "u": ("x", z)},
     )
     ds2 = xr.Dataset(
         data_vars={
@@ -1012,18 +1048,18 @@ def test_merge_dataset(variant, unit, error, dtype):
         coords={
             "x": np.arange(3) * dim_unit,
             "y": np.arange(2, 4) * dim_unit,
-            "z": ("x", np.arange(-3, 0) * coord_unit),
+            "u": ("x", np.arange(-3, 0) * coord_unit),
         },
     )
     ds3 = xr.Dataset(
         data_vars={
-            "a": (("y", "x"), np.zeros_like(array1) * np.nan * data_unit),
-            "b": (("y", "x"), np.zeros_like(array2) * np.nan * data_unit),
+            "a": (("y", "x"), np.full_like(array1, np.nan) * data_unit),
+            "b": (("y", "x"), np.full_like(array2, np.nan) * data_unit),
         },
         coords={
             "x": np.arange(3, 6) * dim_unit,
             "y": np.arange(4, 6) * dim_unit,
-            "z": ("x", np.arange(3, 6) * coord_unit),
+            "u": ("x", np.arange(3, 6) * coord_unit),
         },
     )
 
@@ -1036,8 +1072,19 @@ def test_merge_dataset(variant, unit, error, dtype):
 
     units = extract_units(ds1)
     convert_and_strip = lambda ds: strip_units(convert_units(ds, units))
-    expected = attach_units(
-        func([strip_units(ds1), convert_and_strip(ds2), convert_and_strip(ds3)]), units
+    expected_units = {
+        "a": data_unit,
+        "b": data_unit,
+        "x": original_unit,
+        "y": original_unit,
+        "u": coord_unit,
+    }
+    expected = convert_units(
+        attach_units(
+            func([strip_units(ds1), convert_and_strip(ds2), convert_and_strip(ds3)]),
+            {**units, **{"a": original_unit, "b": original_unit, "z": original_unit}},
+        ),
+        expected_units,
     )
     actual = func([ds1, ds2, ds3])
 
@@ -1051,7 +1098,8 @@ def test_replication_dataarray(func, dtype):
     data_array = xr.DataArray(data=array, dims="x")
 
     numpy_func = getattr(np, func.__name__)
-    expected = xr.DataArray(data=numpy_func(array), dims="x")
+    units = extract_units(numpy_func(data_array))
+    expected = attach_units(func(data_array), units)
     actual = func(data_array)
 
     assert_identical(expected, actual)
@@ -1071,9 +1119,9 @@ def test_replication_dataset(func, dtype):
     )
 
     numpy_func = getattr(np, func.__name__)
-    expected = ds.copy(
-        data={name: numpy_func(array.data) for name, array in ds.data_vars.items()}
-    )
+    units = extract_units(ds.map(numpy_func))
+    expected = attach_units(func(strip_units(ds)), units)
+
     actual = func(ds)
 
     assert_identical(expected, actual)

From 95c3fc1b2e18089da371acbc01db87b5be45036b Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Thu, 12 Dec 2019 11:31:53 +0100
Subject: [PATCH 08/19] remove some xfails

---
 xarray/tests/test_units.py | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index cb01ec049e4..845e3e95852 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -693,12 +693,7 @@ def test_combine_by_coords(variant, unit, error, dtype):
             unit_registry.dimensionless, DimensionalityError, id="dimensionless"
         ),
         pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
-        pytest.param(
-            unit_registry.mm,
-            None,
-            id="compatible_unit",
-            marks=pytest.mark.xfail(reason="wrong order of arguments to `where`"),
-        ),
+        pytest.param(unit_registry.mm, None, id="compatible_unit"),
         pytest.param(unit_registry.m, None, id="identical_unit"),
     ),
     ids=repr,
@@ -891,7 +886,6 @@ def test_concat_dataset(variant, unit, error, dtype):
     assert_identical(expected, actual)
 
 
-@pytest.mark.xfail(reason="blocked by `reindex` / `where`")
 @pytest.mark.parametrize(
     "unit,error",
     (
@@ -997,7 +991,6 @@ def test_merge_dataarray(variant, unit, error, dtype):
     assert_allclose(expected, actual)
 
 
-@pytest.mark.xfail(reason="blocked by `reindex` / `where`")
 @pytest.mark.parametrize(
     "unit,error",
     (

From 81c16dbdf16da58443f0cbcbf9c075ffd2dba60e Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Thu, 12 Dec 2019 11:41:54 +0100
Subject: [PATCH 09/19] use the unit registry's quantity class

---
 xarray/tests/test_units.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 845e3e95852..4bb20212e4e 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -13,13 +13,13 @@
 pint = pytest.importorskip("pint")
 DimensionalityError = pint.errors.DimensionalityError
 
-with warnings.catch_warnings():
-    warnings.simplefilter("ignore")
-    pint.Quantity([])
-
 unit_registry = pint.UnitRegistry()
 Quantity = unit_registry.Quantity
 
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+    Quantity([])
+
 pytestmark = [
     pytest.mark.skipif(
         not IS_NEP18_ACTIVE, reason="NUMPY_EXPERIMENTAL_ARRAY_FUNCTION is not enabled"

From e989406e4cef5e11c3febddfa653cdb6abcbe02e Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Thu, 12 Dec 2019 17:34:08 +0100
Subject: [PATCH 10/19] explain the catch_warnings block

---
 xarray/tests/test_units.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 4bb20212e4e..6675f06c16a 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -16,6 +16,7 @@
 unit_registry = pint.UnitRegistry()
 Quantity = unit_registry.Quantity
 
+# silence pint's BehaviorChangeWarning
 with warnings.catch_warnings():
     warnings.simplefilter("ignore")
     Quantity([])

From d2459e6383348e403288ae31d2bd5b35ce374bd3 Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Thu, 12 Dec 2019 17:34:57 +0100
Subject: [PATCH 11/19] don't use the function wrapper class if we don't need
 arguments

---
 xarray/tests/test_units.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 6675f06c16a..f51beb43cec 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -958,10 +958,9 @@ def test_merge_dataarray(variant, unit, error, dtype):
         dims=("y", "z"),
     )
 
-    func = function(xr.merge)
     if error is not None and variant != "data":
         with pytest.raises(error):
-            func([arr1, arr2, arr3])
+            xr.merge([arr1, arr2, arr3])
 
         return
 
@@ -980,13 +979,15 @@ def test_merge_dataarray(variant, unit, error, dtype):
     }
     expected = convert_units(
         attach_units(
-            func([strip_units(arr1), convert_and_strip(arr2), convert_and_strip(arr3)]),
+            xr.merge(
+                [strip_units(arr1), convert_and_strip(arr2), convert_and_strip(arr3)]
+            ),
             {**units, **{"a": original_unit, "b": data_unit, "c": data_unit}},
         ),
         expected_units,
     )
 
-    actual = func([arr1, arr2, arr3])
+    actual = xr.merge([arr1, arr2, arr3])
 
     assert extract_units(expected) == extract_units(actual)
     assert_allclose(expected, actual)

From 3d3b4818c6d29eda3e09818b420eb2aab6d598a8 Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Thu, 19 Dec 2019 01:30:50 +0100
Subject: [PATCH 12/19] whats-new.rst

---
 doc/whats-new.rst | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index fe05a4d2c21..0d31e280c59 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -35,6 +35,8 @@ New Features
   and support `.dt` accessor for timedelta 
   via :py:class:`core.accessor_dt.TimedeltaAccessor` (:pull:`3612`)
   By `Anderson Banihirwe <https://github.com/andersy005>`_.
+- Support unit aware arrays with pint (:pull:`3611`)
+  By `Justus Magin <https://github.com/keewis>`_.
 
 Bug fixes
 ~~~~~~~~~

From a2f7bca4b9d947caa26e0f2c72bedbef9a3c23e8 Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Wed, 8 Jan 2020 00:08:51 +0100
Subject: [PATCH 13/19] require the new pint version

---
 ci/requirements/py36-min-nep18.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ci/requirements/py36-min-nep18.yml b/ci/requirements/py36-min-nep18.yml
index fc9523ce249..a4be6a82595 100644
--- a/ci/requirements/py36-min-nep18.yml
+++ b/ci/requirements/py36-min-nep18.yml
@@ -10,7 +10,7 @@ dependencies:
   - distributed=2.4
   - numpy=1.17
   - pandas=0.24
-  - pint=0.9  # Actually not enough as it doesn't implement __array_function__yet!
+  - pint=0.10
   - pytest
   - pytest-cov
   - pytest-env

From 91a7cdbd2f169140a0e666ff15dc16cc8853e2cf Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Wed, 15 Jan 2020 19:40:40 +0100
Subject: [PATCH 14/19] use functools.partial instead of function

---
 xarray/tests/test_units.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index cfa1107c12d..48e6ffbe3b3 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -1,3 +1,4 @@
+import functools
 import operator
 import warnings
 
@@ -422,9 +423,8 @@ def __repr__(self):
         return f"function_{self.name}"
 
 
-@pytest.mark.xfail(reason="test bug: apply_ufunc should not be called that way")
 def test_apply_ufunc_dataarray(dtype):
-    func = function(
+    func = functools.partial(
         xr.apply_ufunc, np.mean, input_core_dims=[["x"]], kwargs={"axis": -1}
     )
 
@@ -438,9 +438,8 @@ def test_apply_ufunc_dataarray(dtype):
     assert_identical(expected, actual)
 
 
-@pytest.mark.xfail(reason="test bug: apply_ufunc should not be called that way")
 def test_apply_ufunc_dataset(dtype):
-    func = function(
+    func = functools.partial(
         xr.apply_ufunc, np.mean, input_core_dims=[["x"]], kwargs={"axis": -1}
     )
 

From bfb0bb84cd1e6e14d7656b21a7e9521041ba31bd Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Thu, 16 Jan 2020 02:21:33 +0100
Subject: [PATCH 15/19] remove the convert_from parameter of array_attach_units

---
 xarray/tests/test_units.py | 52 ++++++--------------------------------
 1 file changed, 8 insertions(+), 44 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 48e6ffbe3b3..4fbf77c3d40 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -74,53 +74,17 @@ def array_strip_units(array):
         return array
 
 
-def array_attach_units(data, unit, convert_from=None):
-    try:
-        unit, convert_from = unit
-    except TypeError:
-        pass
-
+def array_attach_units(data, unit):
     if isinstance(data, Quantity):
-        if not convert_from:
-            raise ValueError(
-                "cannot attach unit {unit} to quantity ({data.units})".format(
-                    unit=unit, data=data
-                )
-            )
-        elif isinstance(convert_from, unit_registry.Unit):
-            data = data.magnitude
-        elif convert_from is True:  # intentionally accept exactly true
-            if data.check(unit):
-                convert_from = data.units
-                data = data.magnitude
-            else:
-                raise ValueError(
-                    "cannot convert quantity ({data.units}) to {unit}".format(
-                        unit=unit, data=data
-                    )
-                )
-        else:
-            raise ValueError(
-                "cannot convert from invalid unit {convert_from}".format(
-                    convert_from=convert_from
-                )
-            )
+        raise ValueError(f"cannot attach unit {unit} to quantity {data}")
 
-    # to make sure we also encounter the case of "equal if converted"
-    if convert_from is not None:
-        quantity = (data * convert_from).to(
-            unit
-            if isinstance(unit, unit_registry.Unit)
-            else unit_registry.dimensionless
-        )
-    else:
-        try:
-            quantity = data * unit
-        except np.core._exceptions.UFuncTypeError:
-            if unit != 1:
-                raise
+    try:
+        quantity = data * unit
+    except np.core._exceptions.UFuncTypeError:
+        if isinstance(unit, unit_registry.Unit):
+            raise
 
-            quantity = data
+        quantity = data
 
     return quantity
 

From 853d9e3cf67669e8a3a705bb398d87e52ba67ffd Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Fri, 17 Jan 2020 02:07:36 +0100
Subject: [PATCH 16/19] make sure every top-level function test uses
 assert_units_equal

---
 xarray/tests/test_units.py | 39 +++++++++++++++++++++++++++-----------
 1 file changed, 28 insertions(+), 11 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 4fbf77c3d40..d022d764616 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -399,6 +399,7 @@ def test_apply_ufunc_dataarray(dtype):
     expected = attach_units(func(strip_units(data_array)), extract_units(data_array))
     actual = func(data_array)
 
+    assert_units_equal(expected, actual)
     assert_identical(expected, actual)
 
 
@@ -421,6 +422,7 @@ def test_apply_ufunc_dataset(dtype):
     expected = attach_units(func(strip_units(ds)), extract_units(ds))
     actual = func(ds)
 
+    assert_units_equal(expected, actual)
     assert_identical(expected, actual)
 
 
@@ -505,10 +507,10 @@ def test_align_dataarray(fill_value, variant, unit, error, dtype):
 
     actual_a, actual_b = func(data_array1, data_array2)
 
+    assert_units_equal(expected_a, actual_a)
     assert_allclose(expected_a, actual_a)
-    assert extract_units(expected_a) == extract_units(actual_a)
+    assert_units_equal(expected_b, actual_b)
     assert_allclose(expected_b, actual_b)
-    assert extract_units(expected_b) == extract_units(actual_b)
 
 
 @pytest.mark.parametrize(
@@ -591,10 +593,10 @@ def test_align_dataset(fill_value, unit, variant, error, dtype):
 
     actual_a, actual_b = func(ds1, ds2)
 
+    assert_units_equal(expected_a, actual_a)
     assert_allclose(expected_a, actual_a)
-    assert extract_units(expected_a) == extract_units(actual_a)
+    assert_units_equal(expected_b, actual_b)
     assert_allclose(expected_b, actual_b)
-    assert extract_units(expected_b) == extract_units(actual_b)
 
 
 def test_broadcast_dataarray(dtype):
@@ -612,7 +614,9 @@ def test_broadcast_dataarray(dtype):
 
     actual_a, actual_b = xr.broadcast(a, b)
 
+    assert_units_equal(expected_a, actual_a)
     assert_identical(expected_a, actual_a)
+    assert_units_equal(expected_b, actual_b)
     assert_identical(expected_b, actual_b)
 
 
@@ -645,7 +649,9 @@ def test_broadcast_dataset(dtype):
 
     actual_a, actual_b = xr.broadcast(ds, other)
 
+    assert_units_equal(expected_a, actual_a)
     assert_identical(expected_a, actual_a)
+    assert_units_equal(expected_b, actual_b)
     assert_identical(expected_b, actual_b)
 
 
@@ -716,6 +722,7 @@ def test_combine_by_coords(variant, unit, error, dtype):
     )
     actual = xr.combine_by_coords([ds, other])
 
+    assert_units_equal(expected, actual)
     assert_identical(expected, actual)
 
 
@@ -815,6 +822,7 @@ def test_combine_nested(variant, unit, error, dtype):
     )
     actual = func([[ds1, ds2], [ds3, ds4]])
 
+    assert_units_equal(expected, actual)
     assert_identical(expected, actual)
 
 
@@ -867,6 +875,7 @@ def test_concat_dataarray(variant, unit, error, dtype):
     )
     actual = xr.concat([arr1, arr2], dim="x")
 
+    assert_units_equal(expected, actual)
     assert_identical(expected, actual)
 
 
@@ -917,6 +926,7 @@ def test_concat_dataset(variant, unit, error, dtype):
     )
     actual = xr.concat([ds1, ds2], dim="x")
 
+    assert_units_equal(expected, actual)
     assert_identical(expected, actual)
 
 
@@ -1022,7 +1032,7 @@ def test_merge_dataarray(variant, unit, error, dtype):
 
     actual = xr.merge([arr1, arr2, arr3])
 
-    assert extract_units(expected) == extract_units(actual)
+    assert_units_equal(expected, actual)
     assert_allclose(expected, actual)
 
 
@@ -1116,7 +1126,7 @@ def test_merge_dataset(variant, unit, error, dtype):
     )
     actual = func([ds1, ds2, ds3])
 
-    assert extract_units(expected) == extract_units(actual)
+    assert_units_equal(expected, actual)
     assert_allclose(expected, actual)
 
 
@@ -1130,6 +1140,7 @@ def test_replication_dataarray(func, dtype):
     expected = attach_units(func(data_array), units)
     actual = func(data_array)
 
+    assert_units_equal(expected, actual)
     assert_identical(expected, actual)
 
 
@@ -1152,6 +1163,7 @@ def test_replication_dataset(func, dtype):
 
     actual = func(ds)
 
+    assert_units_equal(expected, actual)
     assert_identical(expected, actual)
 
 
@@ -1191,7 +1203,8 @@ def test_replication_full_like_dataarray(unit, error, dtype):
     )
     actual = xr.full_like(data_array, fill_value=fill_value)
 
-    assert_equal_with_units(expected, actual)
+    assert_units_equal(expected, actual)
+    assert_identical(expected, actual)
 
 
 @pytest.mark.xfail(
@@ -1241,7 +1254,8 @@ def test_replication_full_like_dataset(unit, error, dtype):
     )
     actual = xr.full_like(ds, fill_value=fill_value)
 
-    assert_equal_with_units(expected, actual)
+    assert_units_equal(expected, actual)
+    assert_identical(expected, actual)
 
 
 @pytest.mark.parametrize(
@@ -1283,7 +1297,8 @@ def test_where_dataarray(fill_value, unit, error, dtype):
     )
     actual = xr.where(cond, x, fill_value)
 
-    assert_equal_with_units(expected, actual)
+    assert_units_equal(expected, actual)
+    assert_identical(expected, actual)
 
 
 @pytest.mark.parametrize(
@@ -1327,7 +1342,8 @@ def test_where_dataset(fill_value, unit, error, dtype):
     )
     actual = xr.where(cond, ds, fill_value)
 
-    assert_equal_with_units(expected, actual)
+    assert_units_equal(expected, actual)
+    assert_identical(expected, actual)
 
 
 def test_dot_dataarray(dtype):
@@ -1348,7 +1364,8 @@ def test_dot_dataarray(dtype):
     )
     actual = xr.dot(data_array, other)
 
-    assert_equal_with_units(expected, actual)
+    assert_units_equal(expected, actual)
+    assert_identical(expected, actual)
 
 
 def delete_attrs(*to_delete):

From 94f4a32d86ca6bf39980fbf8e94f6372eed14bef Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Fri, 17 Jan 2020 13:27:13 +0100
Subject: [PATCH 17/19] hide the traceback of the unit comparison function

---
 xarray/tests/test_units.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index d022d764616..9d534e91170 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -255,6 +255,7 @@ def convert_units(obj, to):
 
 
 def assert_units_equal(a, b):
+    __tracebackhide__ = True
     assert extract_units(a) == extract_units(b)
 
 

From 36769816fa9b5112bb879aa47c37715c150a1848 Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Sat, 7 Mar 2020 01:14:49 +0100
Subject: [PATCH 18/19] considerably simplify the merge_dataarray test

---
 xarray/tests/test_units.py | 65 +++++++++++++++++++-------------------
 1 file changed, 33 insertions(+), 32 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 55542f3b371..ac9b8a1c217 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -961,70 +961,71 @@ def test_merge_dataarray(variant, unit, error, dtype):
     data_unit, dim_unit, coord_unit = variants.get(variant)
 
     array1 = np.linspace(0, 1, 2 * 3).reshape(2, 3).astype(dtype) * original_unit
+    x1 = np.arange(2) * original_unit
+    y1 = np.arange(3) * original_unit
+    u1 = np.linspace(10, 20, 2) * original_unit
+    v1 = np.linspace(10, 20, 3) * original_unit
+
     array2 = np.linspace(1, 2, 2 * 4).reshape(2, 4).astype(dtype) * data_unit
-    array3 = np.linspace(0, 2, 3 * 4).reshape(3, 4).astype(dtype) * data_unit
+    x2 = np.arange(2, 4) * dim_unit
+    z2 = np.arange(4) * original_unit
+    u2 = np.linspace(20, 30, 2) * coord_unit
+    w2 = np.linspace(10, 20, 4) * original_unit
 
-    x = np.arange(2) * original_unit
-    y = np.arange(3) * original_unit
-    z = np.arange(4) * original_unit
-    u = np.linspace(10, 20, 2) * original_unit
-    v = np.linspace(10, 20, 3) * original_unit
-    w = np.linspace(10, 20, 4) * original_unit
+    array3 = np.linspace(0, 2, 3 * 4).reshape(3, 4).astype(dtype) * data_unit
+    y3 = np.arange(3, 6) * dim_unit
+    z3 = np.arange(4, 8) * dim_unit
+    v3 = np.linspace(10, 20, 3) * coord_unit
+    w3 = np.linspace(10, 20, 4) * coord_unit
 
     arr1 = xr.DataArray(
         name="a",
         data=array1,
-        coords={"x": x, "y": y, "u": ("x", u), "v": ("y", v)},
+        coords={"x": x1, "y": y1, "u": ("x", u1), "v": ("y", v1)},
         dims=("x", "y"),
     )
     arr2 = xr.DataArray(
-        name="b",
+        name="a",
         data=array2,
-        coords={
-            "x": np.arange(2, 4) * dim_unit,
-            "z": z,
-            "u": ("x", np.linspace(20, 30, 2) * coord_unit),
-            "w": ("z", w),
-        },
+        coords={"x": x2, "z": z2, "u": ("x", u2), "w": ("z", w2)},
         dims=("x", "z"),
     )
     arr3 = xr.DataArray(
-        name="c",
+        name="a",
         data=array3,
-        coords={
-            "y": np.arange(3, 6) * dim_unit,
-            "z": np.arange(4, 8) * dim_unit,
-            "v": ("y", np.linspace(10, 20, 3) * coord_unit),
-            "w": ("z", np.linspace(10, 20, 4) * coord_unit),
-        },
+        coords={"y": y3, "z": z3, "v": ("y", v3), "w": ("z", w3)},
         dims=("y", "z"),
     )
 
-    if error is not None and variant != "data":
+    if error is not None:
         with pytest.raises(error):
             xr.merge([arr1, arr2, arr3])
 
         return
 
-    units = {name: original_unit for name in list("uvwxyz")}
+    units = {name: original_unit for name in list("axyzuvw")}
+
     convert_and_strip = lambda arr: strip_units(convert_units(arr, units))
     expected_units = {
         "a": original_unit,
-        "b": data_unit,
-        "c": data_unit,
-        "u": coord_unit,
-        "v": coord_unit,
-        "w": coord_unit,
+        "u": original_unit,
+        "v": original_unit,
+        "w": original_unit,
         "x": original_unit,
         "y": original_unit,
-        "z": dim_unit,
+        "z": original_unit,
     }
+
     expected = convert_units(
         attach_units(
             xr.merge(
-                [strip_units(arr1), convert_and_strip(arr2), convert_and_strip(arr3)]
+                [
+                    convert_and_strip(arr1),
+                    convert_and_strip(arr2),
+                    convert_and_strip(arr3),
+                ]
             ),
-            {**units, **{"a": original_unit, "b": data_unit, "c": data_unit}},
+            units,
         ),
         expected_units,
     )

From 2eff81db0996ebcce546e8df08d619858e1c580f Mon Sep 17 00:00:00 2001
From: Keewis <keewis@posteo.de>
Date: Sat, 7 Mar 2020 01:20:14 +0100
Subject: [PATCH 19/19] simplify the merge_dataset test

---
 xarray/tests/test_units.py | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index ac9b8a1c217..5577ad7911f 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -1110,17 +1110,13 @@ def test_merge_dataset(variant, unit, error, dtype):
 
     units = extract_units(ds1)
     convert_and_strip = lambda ds: strip_units(convert_units(ds, units))
-    expected_units = {
-        "a": data_unit,
-        "b": data_unit,
-        "x": original_unit,
-        "y": original_unit,
-        "u": coord_unit,
-    }
+    expected_units = {name: original_unit for name in list("abxyzu")}
     expected = convert_units(
         attach_units(
-            func([strip_units(ds1), convert_and_strip(ds2), convert_and_strip(ds3)]),
-            {**units, **{"a": original_unit, "b": original_unit, "z": original_unit}},
+            func(
+                [convert_and_strip(ds1), convert_and_strip(ds2), convert_and_strip(ds3)]
+            ),
+            units,
         ),
         expected_units,
     )