From bbb14a5d8383520f1a1e7e6d885c03ecddfbcf47 Mon Sep 17 00:00:00 2001 From: Fabien Maussion Date: Mon, 9 May 2022 17:25:01 +0200 Subject: [PATCH] Allow string formatting of scalar DataArrays (#5981) * Allow string formatting of scalar DataArrays * better comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * forgot type check * yeah, typing is new to me * Simpler: pass to numpy in all cases * Add dask test * [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> --- xarray/core/common.py | 4 ++++ xarray/tests/test_formatting.py | 22 +++++++++++++++++++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index 3db9b1cfa0c..cf02bcff77b 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -158,6 +158,10 @@ def _repr_html_(self): return f"
{escape(repr(self))}
" return formatting_html.array_repr(self) + def __format__(self: Any, format_spec: str) -> str: + # we use numpy: scalars will print fine and arrays will raise + return self.values.__format__(format_spec) + def _iter(self: Any) -> Iterator[Any]: for n in range(len(self)): yield self[n] diff --git a/xarray/tests/test_formatting.py b/xarray/tests/test_formatting.py index 4bbf41c7b38..a5c044d8ea7 100644 --- a/xarray/tests/test_formatting.py +++ b/xarray/tests/test_formatting.py @@ -9,7 +9,7 @@ import xarray as xr from xarray.core import formatting -from . import requires_netCDF4 +from . import requires_dask, requires_netCDF4 class TestFormatting: @@ -418,6 +418,26 @@ def test_array_repr_variable(self) -> None: with xr.set_options(display_expand_data=False): formatting.array_repr(var) + @requires_dask + def test_array_scalar_format(self) -> None: + var = xr.DataArray(0) + assert var.__format__("") == "0" + assert var.__format__("d") == "0" + assert var.__format__(".2f") == "0.00" + + var = xr.DataArray([0.1, 0.2]) + assert var.__format__("") == "[0.1 0.2]" + with pytest.raises(TypeError) as excinfo: + var.__format__(".2f") + assert "unsupported format string passed to" in str(excinfo.value) + + # also check for dask + var = var.chunk(chunks={"dim_0": 1}) + assert var.__format__("") == "[0.1 0.2]" + with pytest.raises(TypeError) as excinfo: + var.__format__(".2f") + assert "unsupported format string passed to" in str(excinfo.value) + def test_inline_variable_array_repr_custom_repr() -> None: class CustomArray: