Skip to content

Commit

Permalink
Update Datatree html repr to indicate inheritance (#9633)
Browse files Browse the repository at this point in the history
* _inherited_vars -> inherited_vars

* implementation using Coordinates

* datatree.DataTree -> xarray.DataTree

* only show inherited coordinates on root

* test that there is an Inherited coordinates header
  • Loading branch information
TomNicholas authored Oct 21, 2024
1 parent e58edcc commit cfaa72f
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 10 deletions.
6 changes: 3 additions & 3 deletions xarray/core/formatting.py
Original file line number Diff line number Diff line change
Expand Up @@ -447,7 +447,7 @@ def coords_repr(coords: AbstractCoordinates, col_width=None, max_rows=None):


def inherited_coords_repr(node: DataTree, col_width=None, max_rows=None):
coords = _inherited_vars(node._coord_variables)
coords = inherited_vars(node._coord_variables)
if col_width is None:
col_width = _calculate_col_width(coords)
return _mapping_repr(
Expand Down Expand Up @@ -1075,7 +1075,7 @@ def diff_datatree_repr(a: DataTree, b: DataTree, compat):
return "\n\n".join(summary)


def _inherited_vars(mapping: ChainMap) -> dict:
def inherited_vars(mapping: ChainMap) -> dict:
return {k: v for k, v in mapping.parents.items() if k not in mapping.maps[0]}


Expand All @@ -1085,7 +1085,7 @@ def _datatree_node_repr(node: DataTree, show_inherited: bool) -> str:
col_width = _calculate_col_width(node.variables)
max_rows = OPTIONS["display_max_rows"]

inherited_coords = _inherited_vars(node._coord_variables)
inherited_coords = inherited_vars(node._coord_variables)

# Only show dimensions if also showing a variable or coordinates section.
show_dims = (
Expand Down
38 changes: 31 additions & 7 deletions xarray/core/formatting_html.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from typing import TYPE_CHECKING

from xarray.core.formatting import (
inherited_vars,
inline_index_repr,
inline_variable_array_repr,
short_data_repr,
Expand Down Expand Up @@ -248,7 +249,6 @@ def array_section(obj) -> str:
expand_option_name="display_expand_coords",
)


datavar_section = partial(
_mapping_section,
name="Data variables",
Expand Down Expand Up @@ -382,16 +382,40 @@ def summarize_datatree_children(children: Mapping[str, DataTree]) -> str:
expand_option_name="display_expand_groups",
)

inherited_coord_section = partial(
_mapping_section,
name="Inherited coordinates",
details_func=summarize_coords,
max_items_collapse=25,
expand_option_name="display_expand_coords",
)


def datatree_node_repr(group_title: str, node: DataTree, show_inherited=False) -> str:
from xarray.core.coordinates import Coordinates

def datatree_node_repr(group_title: str, dt: DataTree) -> str:
header_components = [f"<div class='xr-obj-type'>{escape(group_title)}</div>"]

ds = dt._to_dataset_view(rebuild_dims=False, inherit=True)
ds = node._to_dataset_view(rebuild_dims=False, inherit=True)
node_coords = node.to_dataset(inherit=False).coords

# use this class to get access to .xindexes property
inherited_coords = Coordinates(
coords=inherited_vars(node._coord_variables),
indexes=inherited_vars(node._indexes),
)

sections = [
children_section(dt.children),
children_section(node.children),
dim_section(ds),
coord_section(ds.coords),
coord_section(node_coords),
]

# only show inherited coordinates on the root
if show_inherited:
sections.append(inherited_coord_section(inherited_coords))

sections += [
datavar_section(ds.data_vars),
attr_section(ds.attrs),
]
Expand Down Expand Up @@ -470,5 +494,5 @@ def _wrap_datatree_repr(r: str, end: bool = False) -> str:


def datatree_repr(dt: DataTree) -> str:
obj_type = f"datatree.{type(dt).__name__}"
return datatree_node_repr(obj_type, dt)
obj_type = f"xarray.{type(dt).__name__}"
return datatree_node_repr(obj_type, dt, show_inherited=True)
19 changes: 19 additions & 0 deletions xarray/tests/test_formatting_html.py
Original file line number Diff line number Diff line change
Expand Up @@ -320,6 +320,25 @@ def test_two_children(
)


class TestDataTreeInheritance:
def test_inherited_section_present(self) -> None:
dt = xr.DataTree.from_dict(
{
"/": None,
"a": None,
}
)
with xr.set_options(display_style="html"):
html = dt._repr_html_().strip()
# checks that the section appears somewhere
assert "Inherited coordinates" in html

# TODO how can we assert that the Inherited coordinates section does not appear in the child group?
# with xr.set_options(display_style="html"):
# child_html = dt["a"]._repr_html_().strip()
# assert "Inherited coordinates" not in child_html


class Test__wrap_datatree_repr:
"""
Unit tests for _wrap_datatree_repr.
Expand Down

0 comments on commit cfaa72f

Please sign in to comment.