From cfaa72f0e9ca86873c10f8bbbd8a62f97fd13f03 Mon Sep 17 00:00:00 2001
From: Tom Nicholas <tom@cworthy.org>
Date: Mon, 21 Oct 2024 09:52:26 -0600
Subject: [PATCH] Update Datatree html repr to indicate inheritance (#9633)

* _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
---
 xarray/core/formatting.py            |  6 ++---
 xarray/core/formatting_html.py       | 38 +++++++++++++++++++++++-----
 xarray/tests/test_formatting_html.py | 19 ++++++++++++++
 3 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/xarray/core/formatting.py b/xarray/core/formatting.py
index e7f2a80caa5..3e23e41ff93 100644
--- a/xarray/core/formatting.py
+++ b/xarray/core/formatting.py
@@ -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(
@@ -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]}
 
 
@@ -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 = (
diff --git a/xarray/core/formatting_html.py b/xarray/core/formatting_html.py
index e0a461caea7..d0cb7c30e91 100644
--- a/xarray/core/formatting_html.py
+++ b/xarray/core/formatting_html.py
@@ -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,
@@ -248,7 +249,6 @@ def array_section(obj) -> str:
     expand_option_name="display_expand_coords",
 )
 
-
 datavar_section = partial(
     _mapping_section,
     name="Data variables",
@@ -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),
     ]
@@ -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)
diff --git a/xarray/tests/test_formatting_html.py b/xarray/tests/test_formatting_html.py
index 9e325b1258d..b518e7e95d9 100644
--- a/xarray/tests/test_formatting_html.py
+++ b/xarray/tests/test_formatting_html.py
@@ -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.