WIP: testing.assert_* check dtype

xarray/core/dataarray.py

@@ -2907,32 +2907,40 @@ def from_iris(cls, cube: "iris_Cube") -> "DataArray":
 
         return from_iris(cube)
 
-    def _all_compat(self, other: "DataArray", compat_str: str) -> bool:
+    def _all_compat(
+        self, other: "DataArray", compat_str: str, check_dtype: bool = False
+    ) -> bool:
         """Helper function for equals, broadcast_equals, and identical"""
 
         def compat(x, y):
-            return getattr(x.variable, compat_str)(y.variable)
+            return getattr(x.variable, compat_str)(y.variable, check_dtype=check_dtype)
 
         return utils.dict_equiv(self.coords, other.coords, compat=compat) and compat(
             self, other
         )
 
-    def broadcast_equals(self, other: "DataArray") -> bool:
+    def broadcast_equals(self, other: "DataArray", check_dtype: bool = False) -> bool:
         """Two DataArrays are broadcast equal if they are equal after
         broadcasting them against each other such that they have the same
         dimensions.
 
+        Parameters
+        ----------
+        check_dtype : bool, default: False
+           Whether to check if the objects' dtypes are identical. Compares the
+           dtypes of the data and the coords.
+
         See Also
         --------
         DataArray.equals
         DataArray.identical
         """
         try:
-            return self._all_compat(other, "broadcast_equals")
+            return self._all_compat(other, "broadcast_equals", check_dtype=check_dtype)
         except (TypeError, AttributeError):
             return False
 
-    def equals(self, other: "DataArray") -> bool:
+    def equals(self, other: "DataArray", check_dtype: bool = False) -> bool:
         """True if two DataArrays have the same dimensions, coordinates and
         values; otherwise False.
 
@@ -2942,27 +2950,41 @@ def equals(self, other: "DataArray") -> bool:
         This method is necessary because `v1 == v2` for ``DataArray``
         does element-wise comparisons (like numpy.ndarrays).
 
+        Parameters
+        ----------
+        check_dtype : bool, default: False
+           Whether to check if the objects' dtypes are identical. Compares the
+           dtypes of the data and the coords.
+
         See Also
         --------
         DataArray.broadcast_equals
         DataArray.identical
         """
         try:
-            return self._all_compat(other, "equals")
+            return self._all_compat(other, "equals", check_dtype=check_dtype)
         except (TypeError, AttributeError):
             return False
 
-    def identical(self, other: "DataArray") -> bool:
+    def identical(self, other: "DataArray", check_dtype: bool = False) -> bool:
         """Like equals, but also checks the array name and attributes, and
         attributes on all coordinates.
 
+        Parameters
+        ----------
+        check_dtype : bool, default: False
+           Whether to check if the objects' dtypes are identical. Compares the
+           dtypes of the data and the coords.
+
         See Also
         --------
         DataArray.broadcast_equals
         DataArray.equals
         """
         try:
-            return self.name == other.name and self._all_compat(other, "identical")
+            return self.name == other.name and self._all_compat(
+                other, "identical", check_dtype=check_dtype
+            )
         except (TypeError, AttributeError):
             return False
 

xarray/core/dataset.py

@@ -1537,37 +1537,43 @@ def __delitem__(self, key: Hashable) -> None:
     # https://github.com/python/mypy/issues/4266
     __hash__ = None  # type: ignore[assignment]
 
-    def _all_compat(self, other: "Dataset", compat_str: str) -> bool:
+    def _all_compat(self, other: "Dataset", compat_str: str, check_dtype: bool) -> bool:
         """Helper function for equals and identical"""
 
         # some stores (e.g., scipy) do not seem to preserve order, so don't
         # require matching order for equality
         def compat(x: Variable, y: Variable) -> bool:
-            return getattr(x, compat_str)(y)
+            return getattr(x, compat_str)(y, check_dtype=check_dtype)
 
         return self._coord_names == other._coord_names and utils.dict_equiv(
             self._variables, other._variables, compat=compat
         )
 
-    def broadcast_equals(self, other: "Dataset") -> bool:
+    def broadcast_equals(self, other: "Dataset", check_dtype: bool = False) -> bool:
         """Two Datasets are broadcast equal if they are equal after
         broadcasting all variables against each other.
 
         For example, variables that are scalar in one dataset but non-scalar in
         the other dataset can still be broadcast equal if the the non-scalar
         variable is a constant.
 
+        Parameters
+        ----------
+        check_dtype : bool, default: False
+           Whether to check if the objects' dtypes are identical. Compares the
+           dtypes of all data variables and coords.
+
         See Also
         --------
         Dataset.equals
         Dataset.identical
         """
         try:
-            return self._all_compat(other, "broadcast_equals")
+            return self._all_compat(other, "broadcast_equals", check_dtype=check_dtype)
         except (TypeError, AttributeError):
             return False
 
-    def equals(self, other: "Dataset") -> bool:
+    def equals(self, other: "Dataset", check_dtype: bool = False) -> bool:
         """Two Datasets are equal if they have matching variables and
         coordinates, all of which are equal.
 
@@ -1577,28 +1583,40 @@ def equals(self, other: "Dataset") -> bool:
         This method is necessary because `v1 == v2` for ``Dataset``
         does element-wise comparisons (like numpy.ndarrays).
 
+        Parameters
+        ----------
+        check_dtype : bool, default: False
+           Whether to check if the objects' dtypes are identical. Compares the
+           dtypes of all data variables and coords.
+
         See Also
         --------
         Dataset.broadcast_equals
         Dataset.identical
         """
         try:
-            return self._all_compat(other, "equals")
+            return self._all_compat(other, "equals", check_dtype=check_dtype)
         except (TypeError, AttributeError):
             return False
 
-    def identical(self, other: "Dataset") -> bool:
+    def identical(self, other: "Dataset", check_dtype: bool = False) -> bool:
         """Like equals, but also checks all dataset attributes and the
         attributes on all variables and coordinates.
 
+        Parameters
+        ----------
+        check_dtype : bool, default: False
+           Whether to check if the objects' dtypes are identical. Compares the
+           dtypes of all data variables and coords.
+
         See Also
         --------
         Dataset.broadcast_equals
         Dataset.equals
         """
         try:
             return utils.dict_equiv(self.attrs, other.attrs) and self._all_compat(
-                other, "identical"
+                other, "identical", check_dtype=check_dtype
             )
         except (TypeError, AttributeError):
             return False

xarray/core/duck_array_ops.py

@@ -218,10 +218,11 @@ def as_shared_dtype(scalars_or_arrays):
     return [x.astype(out_type, copy=False) for x in arrays]
 
 
-def lazy_array_equiv(arr1, arr2):
+def lazy_array_equiv(arr1, arr2, check_dtype=False):
     """Like array_equal, but doesn't actually compare values.
     Returns True when arr1, arr2 identical or their dask tokens are equal.
     Returns False when shapes are not equal.
+    Returns False if dtype does not match and check_dtype is True.
     Returns None when equality cannot determined: one or both of arr1, arr2 are numpy arrays;
     or their dask tokens are not equal
     """
@@ -231,6 +232,9 @@ def lazy_array_equiv(arr1, arr2):
     arr2 = asarray(arr2)
     if arr1.shape != arr2.shape:
         return False
+    # "is False" needed -> should not return on None
+    if check_dtype and same_dtype(arr1, arr2, lazy=True) is False:
+        return False
     if dask_array and is_duck_dask_array(arr1) and is_duck_dask_array(arr2):
         # GH3068, GH4221
         if tokenize(arr1) == tokenize(arr2):
@@ -240,26 +244,47 @@ def lazy_array_equiv(arr1, arr2):
     return None
 
 
-def allclose_or_equiv(arr1, arr2, rtol=1e-5, atol=1e-8):
+def same_dtype(arr1, arr2, lazy):
+
+    # object dask arrays can change dtype -> need to compute them
+    if arr1.dtype == object and is_duck_dask_array(arr1):
+        if lazy:
+            return None
+        # arr.compute() can return a scalar -> wrap in an array
+        arr1 = asarray(arr1.compute())
+
+    if arr2.dtype == object and is_duck_dask_array(arr2):
+        if lazy:
+            return None
+        arr2 = asarray(arr2.compute())
+
+    return arr1.dtype == arr2.dtype
+
+
+def allclose_or_equiv(arr1, arr2, rtol=1e-5, atol=1e-8, check_dtype=False):
     """Like np.allclose, but also allows values to be NaN in both arrays"""
     arr1 = asarray(arr1)
     arr2 = asarray(arr2)
 
-    lazy_equiv = lazy_array_equiv(arr1, arr2)
+    lazy_equiv = lazy_array_equiv(arr1, arr2, check_dtype=check_dtype)
     if lazy_equiv is None:
+        if check_dtype and not same_dtype(arr1, arr2, lazy=False):
+            return False
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", r"All-NaN (slice|axis) encountered")
             return bool(isclose(arr1, arr2, rtol=rtol, atol=atol, equal_nan=True).all())
     else:
         return lazy_equiv
 
 
-def array_equiv(arr1, arr2):
+def array_equiv(arr1, arr2, check_dtype=False):
     """Like np.array_equal, but also allows values to be NaN in both arrays"""
     arr1 = asarray(arr1)
     arr2 = asarray(arr2)
-    lazy_equiv = lazy_array_equiv(arr1, arr2)
+    lazy_equiv = lazy_array_equiv(arr1, arr2, check_dtype=check_dtype)
     if lazy_equiv is None:
+        if check_dtype and not same_dtype(arr1, arr2, lazy=False):
+            return False
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", "In the future, 'NAT == x'")
             flag_array = (arr1 == arr2) | (isnull(arr1) & isnull(arr2))
@@ -268,14 +293,16 @@ def array_equiv(arr1, arr2):
         return lazy_equiv
 
 
-def array_notnull_equiv(arr1, arr2):
+def array_notnull_equiv(arr1, arr2, check_dtype=False):
     """Like np.array_equal, but also allows values to be NaN in either or both
     arrays
     """
     arr1 = asarray(arr1)
     arr2 = asarray(arr2)
-    lazy_equiv = lazy_array_equiv(arr1, arr2)
+    lazy_equiv = lazy_array_equiv(arr1, arr2, check_dtype=check_dtype)
     if lazy_equiv is None:
+        if check_dtype and not same_dtype(arr1, arr2, lazy=False):
+            return False
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", "In the future, 'NAT == x'")
             flag_array = (arr1 == arr2) | isnull(arr1) | isnull(arr2)

xarray/core/formatting.py

@@ -566,7 +566,9 @@ def diff_dim_summary(a, b):
         return ""
 
 
-def _diff_mapping_repr(a_mapping, b_mapping, compat, title, summarizer, col_width=None):
+def _diff_mapping_repr(
+    a_mapping, b_mapping, compat, title, summarizer, col_width=None, check_dtype=False
+):
     def extra_items_repr(extra_keys, mapping, ab_side):
         extra_repr = [summarizer(k, mapping[k], col_width) for k in extra_keys]
         if extra_repr:
@@ -586,9 +588,15 @@ def extra_items_repr(extra_keys, mapping, ab_side):
         try:
             # compare xarray variable
             if not callable(compat):
-                compatible = getattr(a_mapping[k], compat)(b_mapping[k])
+                compatible = getattr(a_mapping[k], compat)(
+                    b_mapping[k], check_dtype=check_dtype
+                )
             else:
-                compatible = compat(a_mapping[k], b_mapping[k])
+                compatible = compat(
+                    a_mapping[k],
+                    b_mapping[k],
+                    check_dtype=check_dtype,
+                )
             is_variable = True
         except AttributeError:
             # compare attribute value
@@ -620,8 +628,9 @@ def extra_items_repr(extra_keys, mapping, ab_side):
 
             diff_items += [ab_side + s[1:] for ab_side, s in zip(("L", "R"), temp)]
 
+    maybe_dtype = " (values and/ or dtype)" if check_dtype else ""
     if diff_items:
-        summary += [f"Differing {title.lower()}:"] + diff_items
+        summary += [f"Differing {title.lower()}{maybe_dtype}:"] + diff_items
 
     summary += extra_items_repr(a_keys - b_keys, a_mapping, "left")
     summary += extra_items_repr(b_keys - a_keys, b_mapping, "right")
@@ -656,7 +665,7 @@ def _compat_to_str(compat):
         return compat
 
 
-def diff_array_repr(a, b, compat):
+def diff_array_repr(a, b, compat, check_dtype=False):
     # used for DataArray, Variable and IndexVariable
     summary = [
         "Left and right {} objects are not {}".format(
@@ -670,18 +679,22 @@ def diff_array_repr(a, b, compat):
     else:
         equiv = array_equiv
 
-    if not equiv(a.data, b.data):
+    maybe_dtype = " or dtype" if check_dtype else ""
+
+    if not equiv(a.data, b.data, check_dtype=check_dtype):
         temp = [wrap_indent(short_numpy_repr(obj), start="    ") for obj in (a, b)]
         diff_data_repr = [
             ab_side + "\n" + ab_data_repr
             for ab_side, ab_data_repr in zip(("L", "R"), temp)
         ]
-        summary += ["Differing values:"] + diff_data_repr
+        summary += [f"Differing values{maybe_dtype}:"] + diff_data_repr
 
     if hasattr(a, "coords"):
         col_width = _calculate_col_width(set(a.coords) | set(b.coords))
         summary.append(
-            diff_coords_repr(a.coords, b.coords, compat, col_width=col_width)
+            diff_coords_repr(
+                a.coords, b.coords, compat, col_width=col_width, check_dtype=check_dtype
+            )
         )
 
     if compat == "identical":
@@ -690,7 +703,7 @@ def diff_array_repr(a, b, compat):
     return "\n".join(summary)
 
 
-def diff_dataset_repr(a, b, compat):
+def diff_dataset_repr(a, b, compat, check_dtype=False):
     summary = [
         "Left and right {} objects are not {}".format(
             type(a).__name__, _compat_to_str(compat)
@@ -702,9 +715,19 @@ def diff_dataset_repr(a, b, compat):
     )
 
     summary.append(diff_dim_summary(a, b))
-    summary.append(diff_coords_repr(a.coords, b.coords, compat, col_width=col_width))
     summary.append(
-        diff_data_vars_repr(a.data_vars, b.data_vars, compat, col_width=col_width)
+        diff_coords_repr(
+            a.coords, b.coords, compat, col_width=col_width, check_dtype=check_dtype
+        )
+    )
+    summary.append(
+        diff_data_vars_repr(
+            a.data_vars,
+            b.data_vars,
+            compat,
+            col_width=col_width,
+            check_dtype=check_dtype,
+        )
     )
 
     if compat == "identical":