Skip to content

Display data returned in apply_ufunc error message #8179

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Sep 14, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 13 additions & 12 deletions xarray/core/computation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down
8 changes: 6 additions & 2 deletions xarray/core/formatting.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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:
Expand Down
10 changes: 8 additions & 2 deletions xarray/tests/test_computation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand All @@ -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"):
Expand Down