From ab84e048c2d80c82af692ec52dc37d31fc49f123 Mon Sep 17 00:00:00 2001 From: Stephan Hoyer Date: Sun, 22 Sep 2024 15:56:22 -0700 Subject: [PATCH] Fix DataTree repr to not repeat inherited coordinates (#9532) * Fix DataTree repr to not repeat inherited coordinates Fixes GH9499 * skip failing test on Windows --- xarray/core/formatting.py | 3 +- xarray/tests/test_datatree.py | 107 +++++++++++++++++++++++++++++++++- 2 files changed, 108 insertions(+), 2 deletions(-) diff --git a/xarray/core/formatting.py b/xarray/core/formatting.py index 1cea9a7a28d..646e000a984 100644 --- a/xarray/core/formatting.py +++ b/xarray/core/formatting.py @@ -1102,7 +1102,8 @@ def _datatree_node_repr(node: DataTree, show_inherited: bool) -> str: summary.append(f"{dims_start}({dims_values})") if node._node_coord_variables: - summary.append(coords_repr(node.coords, col_width=col_width, max_rows=max_rows)) + node_coords = node.to_dataset(inherited=False).coords + summary.append(coords_repr(node_coords, col_width=col_width, max_rows=max_rows)) if show_inherited and inherited_coords: summary.append( diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 4e22ba57e2e..2adf73fa5c8 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -1,4 +1,5 @@ import re +import sys import typing from copy import copy, deepcopy from textwrap import dedent @@ -15,6 +16,8 @@ from xarray.testing import assert_equal, assert_identical from xarray.tests import assert_array_equal, create_test_data, source_ndarray +ON_WINDOWS = sys.platform == "win32" + class TestTreeCreation: def test_empty(self): @@ -1052,7 +1055,7 @@ def test_repr_two_children(self): { "/": Dataset(coords={"x": [1.0]}), "/first_child": None, - "/second_child": Dataset({"foo": ("x", [0.0])}), + "/second_child": Dataset({"foo": ("x", [0.0])}, coords={"z": 1.0}), } ) @@ -1067,6 +1070,8 @@ def test_repr_two_children(self): ├── Group: /first_child └── Group: /second_child Dimensions: (x: 1) + Coordinates: + z float64 8B 1.0 Data variables: foo (x) float64 8B 0.0 """ @@ -1091,6 +1096,8 @@ def test_repr_two_children(self): Group: /second_child Dimensions: (x: 1) + Coordinates: + z float64 8B 1.0 Inherited coordinates: * x (x) float64 8B 1.0 Data variables: @@ -1138,6 +1145,104 @@ def test_repr_inherited_dims(self): ).strip() assert result == expected + @pytest.mark.skipif( + ON_WINDOWS, reason="windows (pre NumPy2) uses int32 instead of int64" + ) + def test_doc_example(self): + # regression test for https://github.com/pydata/xarray/issues/9499 + time = xr.DataArray(data=["2022-01", "2023-01"], dims="time") + stations = xr.DataArray(data=list("abcdef"), dims="station") + lon = [-100, -80, -60] + lat = [10, 20, 30] + # Set up fake data + wind_speed = xr.DataArray(np.ones((2, 6)) * 2, dims=("time", "station")) + pressure = xr.DataArray(np.ones((2, 6)) * 3, dims=("time", "station")) + air_temperature = xr.DataArray(np.ones((2, 6)) * 4, dims=("time", "station")) + dewpoint = xr.DataArray(np.ones((2, 6)) * 5, dims=("time", "station")) + infrared = xr.DataArray(np.ones((2, 3, 3)) * 6, dims=("time", "lon", "lat")) + true_color = xr.DataArray(np.ones((2, 3, 3)) * 7, dims=("time", "lon", "lat")) + tree = xr.DataTree.from_dict( + { + "/": xr.Dataset( + coords={"time": time}, + ), + "/weather": xr.Dataset( + coords={"station": stations}, + data_vars={ + "wind_speed": wind_speed, + "pressure": pressure, + }, + ), + "/weather/temperature": xr.Dataset( + data_vars={ + "air_temperature": air_temperature, + "dewpoint": dewpoint, + }, + ), + "/satellite": xr.Dataset( + coords={"lat": lat, "lon": lon}, + data_vars={ + "infrared": infrared, + "true_color": true_color, + }, + ), + }, + ) + + result = repr(tree) + expected = dedent( + """ + + Group: / + │ Dimensions: (time: 2) + │ Coordinates: + │ * time (time) + Group: /weather + │ Dimensions: (time: 2, station: 6) + │ Coordinates: + │ * station (station) str: return re.escape(dedent(message).strip())