Display data returned in apply_ufunc error message (#8179)

* Display data returned in ufunc error message

This makes debugging much easier!

* Use `short_array_repr`

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: max-sixty

xarray/core/computation.py

@@ -659,6 +659,7 @@ def apply_variable_ufunc(
     dask_gufunc_kwargs=None,
 ) -> Variable | tuple[Variable, ...]:
     """Apply a ndarray level function over Variable and/or ndarray objects."""
+    from xarray.core.formatting import short_array_repr
     from xarray.core.variable import Variable, as_compatible_data
 
     dim_sizes = unified_dim_sizes(
@@ -766,11 +767,10 @@ def func(*arrays):
         not isinstance(result_data, tuple) or len(result_data) != signature.num_outputs
     ):
         raise ValueError(
-            "applied function does not have the number of "
-            "outputs specified in the ufunc signature. "
-            "Result is not a tuple of {} elements: {!r}".format(
-                signature.num_outputs, result_data
-            )
+            f"applied function does not have the number of "
+            f"outputs specified in the ufunc signature. "
+            f"Result is not a tuple of {signature.num_outputs} elements:\n\n"
+            f"{short_array_repr(result_data)}"
         )
 
     objs = _all_of_type(args, Variable)
@@ -784,21 +784,22 @@ def func(*arrays):
         data = as_compatible_data(data)
         if data.ndim != len(dims):
             raise ValueError(
-                "applied function returned data with unexpected "
+                "applied function returned data with an unexpected "
                 f"number of dimensions. Received {data.ndim} dimension(s) but "
-                f"expected {len(dims)} dimensions with names: {dims!r}"
+                f"expected {len(dims)} dimensions with names {dims!r}, from:\n\n"
+                f"{short_array_repr(data)}"
             )
 
         var = Variable(dims, data, fastpath=True)
         for dim, new_size in var.sizes.items():
             if dim in dim_sizes and new_size != dim_sizes[dim]:
                 raise ValueError(
-                    "size of dimension {!r} on inputs was unexpectedly "
-                    "changed by applied function from {} to {}. Only "
+                    f"size of dimension '{dim}' on inputs was unexpectedly "
+                    f"changed by applied function from {dim_sizes[dim]} to {new_size}. Only "
                     "dimensions specified in ``exclude_dims`` with "
-                    "xarray.apply_ufunc are allowed to change size.".format(
-                        dim, dim_sizes[dim], new_size
-                    )
+                    "xarray.apply_ufunc are allowed to change size. "
+                    "The data returned was:\n\n"
+                    f"{short_array_repr(data)}"
                 )
 
         var.attrs = attrs

xarray/core/formatting.py

@@ -10,6 +10,7 @@
 from datetime import datetime, timedelta
 from itertools import chain, zip_longest
 from reprlib import recursive_repr
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -21,6 +22,9 @@
 from xarray.core.pycompat import array_type
 from xarray.core.utils import is_duck_array
 
+if TYPE_CHECKING:
+    from xarray.core.coordinates import AbstractCoordinates
+
 
 def pretty_print(x, numchars: int):
     """Given an object `x`, call `str(x)` and format the returned string so
@@ -398,7 +402,7 @@ def _mapping_repr(
 )
 
 
-def coords_repr(coords, col_width=None, max_rows=None):
+def coords_repr(coords: AbstractCoordinates, col_width=None, max_rows=None):
     if col_width is None:
         col_width = _calculate_col_width(coords)
     return _mapping_repr(
@@ -412,7 +416,7 @@ def coords_repr(coords, col_width=None, max_rows=None):
     )
 
 
-def inline_index_repr(index, max_width=None):
+def inline_index_repr(index: pd.Index, max_width=None):
     if hasattr(index, "_repr_inline_"):
         repr_ = index._repr_inline_(max_width=max_width)
     else:

xarray/tests/test_computation.py

@@ -1666,7 +1666,10 @@ def identity(x):
     def tuple3x(x):
         return (x, x, x)
 
-    with pytest.raises(ValueError, match=r"number of outputs"):
+    with pytest.raises(
+        ValueError,
+        match=r"number of outputs.*Result is not a tuple of 2 elements:\n\narray\(\[0",
+    ):
         apply_ufunc(identity, variable, output_core_dims=[(), ()])
 
     with pytest.raises(ValueError, match=r"number of outputs"):
@@ -1682,7 +1685,10 @@ def add_dim(x):
     def remove_dim(x):
         return x[..., 0]
 
-    with pytest.raises(ValueError, match=r"unexpected number of dimensions"):
+    with pytest.raises(
+        ValueError,
+        match=r"unexpected number of dimensions.*from:\n\n.*array\(\[\[0",
+    ):
         apply_ufunc(add_dim, variable, output_core_dims=[("y", "z")])
 
     with pytest.raises(ValueError, match=r"unexpected number of dimensions"):