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>

Author: fmaussion

xarray/core/common.py

@@ -158,6 +158,10 @@ def _repr_html_(self):
             return f"<pre>{escape(repr(self))}</pre>"
         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]

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: