Improve HTML reprs

xarray/core/formatting_html.py

@@ -3,10 +3,10 @@
 import uuid
 from collections import OrderedDict
 from collections.abc import Mapping
+from dataclasses import dataclass
 from functools import lru_cache, partial
 from html import escape
 from importlib.resources import files
-from math import ceil
 from typing import TYPE_CHECKING, Literal
 
 from xarray.core.formatting import (
@@ -172,25 +172,32 @@ def summarize_indexes(indexes) -> str:
 
 
 def collapsible_section(
-    name, inline_details="", details="", n_items=None, enabled=True, collapsed=False
+    header: str,
+    inline_details="",
+    details="",
+    n_items=None,
+    enabled=True,
+    collapsed=False,
 ) -> str:
     # "unique" id to expand/collapse the section
     data_id = "section-" + str(uuid.uuid4())
 
     has_items = n_items is not None and n_items
     n_items_span = "" if n_items is None else f" <span>({n_items})</span>"
-    enabled = "" if enabled and has_items else "disabled"
-    collapsed = "" if collapsed or not has_items else "checked"
+    enabled = "" if enabled and has_items else " disabled"
+    collapsed = "" if collapsed or not has_items else " checked"
     tip = " title='Expand/collapse section'" if enabled else ""
+    span_grid = " xr-span-grid" if not inline_details else ""
 
-    return (
-        f"<input id='{data_id}' class='xr-section-summary-in' "
-        f"type='checkbox' {enabled} {collapsed}>"
-        f"<label for='{data_id}' class='xr-section-summary' {tip}>"
-        f"{name}:{n_items_span}</label>"
-        f"<div class='xr-section-inline-details'>{inline_details}</div>"
-        f"<div class='xr-section-details'>{details}</div>"
+    html = (
+        f"<input id='{data_id}' class='xr-section-summary-in' type='checkbox'{enabled}{collapsed} />"
+        f"<label for='{data_id}' class='xr-section-summary{span_grid}' {tip}>{header}{n_items_span}</label>"
     )
+    if inline_details:
+        html += f"<div class='xr-section-inline-details'>{inline_details}</div>"
+    if details:
+        html += f"<div class='xr-section-details'>{details}</div>"
+    return html
 
 
 def _mapping_section(
@@ -201,9 +208,10 @@ def _mapping_section(
     expand_option_name,
     enabled=True,
     max_option_name: Literal["display_max_children"] | None = None,
+    **kwargs,
 ) -> str:
     n_items = len(mapping)
-    expanded = _get_boolean_with_default(
+    expanded = max_items_collapse is None or _get_boolean_with_default(
         expand_option_name, n_items < max_items_collapse
     )
     collapsed = not expanded
@@ -215,9 +223,9 @@ def _mapping_section(
             inline_details = f"({max_items}/{n_items})"
 
     return collapsible_section(
-        name,
+        f"{name}:",
         inline_details=inline_details,
-        details=details_func(mapping),
+        details=details_func(mapping, **kwargs),
         n_items=n_items,
         enabled=enabled,
         collapsed=collapsed,
@@ -228,7 +236,7 @@ def dim_section(obj) -> str:
     dim_list = format_dims(obj.sizes, obj.xindexes.dims)
 
     return collapsible_section(
-        "Dimensions", inline_details=dim_list, enabled=False, collapsed=True
+        "Dimensions:", inline_details=dim_list, enabled=False, collapsed=True
     )
 
 
@@ -294,14 +302,18 @@ def _get_indexes_dict(indexes):
     }
 
 
+def _sections_repr(sections: list[str]) -> str:
+    section_items = "".join(f"<li class='xr-section-item'>{s}</li>" for s in sections)
+    return f"<ul class='xr-sections'>{section_items}</ul>"
+
+
 def _obj_repr(obj, header_components, sections):
     """Return HTML repr of an xarray object.
 
     If CSS is not injected (untrusted notebook), fallback to the plain text repr.
 
     """
     header = f"<div class='xr-header'>{''.join(h for h in header_components)}</div>"
-    sections = "".join(f"<li class='xr-section-item'>{s}</li>" for s in sections)
 
     icons_svg, css_style = _load_static_files()
     return (
@@ -310,7 +322,7 @@ def _obj_repr(obj, header_components, sections):
         f"<pre class='xr-text-repr-fallback'>{escape(repr(obj))}</pre>"
         "<div class='xr-wrap' style='display:none'>"
         f"{header}"
-        f"<ul class='xr-sections'>{sections}</ul>"
+        f"{_sections_repr(sections)}"
         "</div>"
         "</div>"
     )
@@ -384,7 +396,16 @@ def dataset_repr(ds) -> str:
     return _obj_repr(ds, header_components, sections)
 
 
-def datatree_node_sections(node: DataTree, root: bool = False) -> list[str]:
+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_sections(node: DataTree, root: bool) -> tuple[list[str], int]:
     from xarray.core.coordinates import Coordinates
 
     ds = node._to_dataset_view(rebuild_dims=False, inherit=True)
@@ -397,78 +418,146 @@ def datatree_node_sections(node: DataTree, root: bool = False) -> list[str]:
     )
 
     # Only show dimensions if also showing a variable or coordinates section.
-    show_dims = (
-        node._node_coord_variables
-        or (root and inherited_coords)
-        or node._data_variables
-    )
+    show_dims = node_coords or (root and inherited_coords) or ds.data_vars
 
     sections = []
-
-    if node.children:
-        children_max_items = 1 if ds.data_vars else 6
-        sections.append(
-            children_section(node.children, max_items_collapse=children_max_items)
-        )
-
     if show_dims:
         sections.append(dim_section(ds))
-
     if node_coords:
         sections.append(coord_section(node_coords))
-
-    # only show inherited coordinates on the root
     if root and inherited_coords:
         sections.append(inherited_coord_section(inherited_coords))
-
     if ds.data_vars:
         sections.append(datavar_section(ds.data_vars))
-
     if ds.attrs:
         sections.append(attr_section(ds.attrs))
 
-    return sections
+    displayed_line_count = (
+        len(node.children)
+        + int(bool(show_dims))
+        + int(bool(node_coords))
+        + len(node_coords)
+        + int(root) * (int(bool(inherited_coords)) + len(inherited_coords))
+        + int(bool(ds.data_vars))
+        + len(ds.data_vars)
+        + int(bool(ds.attrs))
+        + len(ds.attrs)
+    )
 
+    return sections, displayed_line_count
 
-def summarize_datatree_children(children: Mapping[str, DataTree]) -> str:
-    MAX_CHILDREN = OPTIONS["display_max_children"]
-    n_children = len(children)
 
-    children_html = []
-    for i, child in enumerate(children.values()):
-        if i < ceil(MAX_CHILDREN / 2) or i >= ceil(n_children - MAX_CHILDREN / 2):
-            is_last = i == (n_children - 1)
-            children_html.append(datatree_child_repr(child, end=is_last))
-        elif n_children > MAX_CHILDREN and i == ceil(MAX_CHILDREN / 2):
-            children_html.append("<div>...</div>")
-
-    return "".join(
-        [
-            "<div style='display: inline-grid; grid-template-columns: 100%; grid-column: 1 / -1'>",
-            "".join(children_html),
-            "</div>",
-        ]
+def _tree_item_count(node: DataTree, cache: dict[int, int]) -> int:
+    if id(node) in cache:
+        return cache[id(node)]
+
+    node_ds = node.to_dataset(inherit=False)
+    node_count = len(node_ds.variables) + len(node_ds.attrs)
+    child_count = sum(
+        _tree_item_count(child, cache) for child in node.children.values()
+    )
+    total = node_count + child_count
+    cache[id(node)] = total
+    return total
+
+
+@dataclass
+class _DataTreeDisplay:
+    node: DataTree
+    sections: list[str]
+    item_count: int
+    collapsed: bool
+    disabled: bool
+
+
+def _build_datatree_displays(tree: DataTree) -> dict[str, _DataTreeDisplay]:
+    displayed_line_count = 0
+    html_line_count = 0
+    displays: dict[str, _DataTreeDisplay] = {}
+    item_count_cache: dict[int, int] = {}
+    root = True
+    collapsed = False
+    disabled = False
+
+    html_limit = OPTIONS["display_max_html_elements"]
+    uncollapsed_limit = OPTIONS["display_max_items"]
+
+    too_many_items_section = collapsible_section(
+        "<em>Too many items to display (display_max_html_elements exceeded)</em>",
+        enabled=False,
+        collapsed=True,
     )
 
+    for node in tree.subtree:  # breadth-first
+        parent = node.parent
+        if parent is not None:
+            parent_display = displays.get(parent.path, None)
+            if parent_display is not None and parent_display.disabled:
+                break  # no need to build display
 
-children_section = partial(
-    _mapping_section,
-    name="Groups",
-    details_func=summarize_datatree_children,
-    max_option_name="display_max_children",
-    expand_option_name="display_expand_groups",
-)
+        item_count = _tree_item_count(node, item_count_cache)
+
+        sections, node_line_count = _datatree_node_sections(node, root)
+        new_displayed_count = displayed_line_count + node_line_count
+        new_html_count = html_line_count + node_line_count
+
+        disabled = not root and (disabled or new_html_count > html_limit)
+        if disabled:
+            sections = [too_many_items_section]
+            collapsed = True
+        else:
+            html_line_count = new_html_count
+
+        collapsed = not root and (collapsed or new_displayed_count > uncollapsed_limit)
+        if not collapsed:
+            displayed_line_count = new_displayed_count
+
+        displays[node.path] = _DataTreeDisplay(
+            node, sections, item_count, collapsed, disabled
+        )
+        root = False
+
+    # If any node is collapsed, ensure its immediate siblings are also collapsed
+    for display in displays.values():
+        if not display.disabled:
+            if any(
+                displays[child.path].collapsed
+                for child in display.node.children.values()
+            ):
+                for child in display.node.children.values():
+                    displays[child.path].collapsed = True
+
+    return displays
 
-inherited_coord_section = partial(
-    _mapping_section,
-    name="Inherited coordinates",
-    details_func=summarize_coords,
-    max_items_collapse=25,
-    expand_option_name="display_expand_coords",
-)
 
+def children_section(
+    children: Mapping[str, DataTree], displays: dict[str, _DataTreeDisplay]
+) -> str:
+    child_elements = []
+    for i, child in enumerate(children.values()):
+        is_last = i == (len(children) - 1)
+        child_elements.append(datatree_child_repr(child, displays, end=is_last))
+
+    children_html = "".join(child_elements)
+    return f"<div class='xr-children'>{children_html}</div>"
+
+
+def datatree_sections(
+    node: DataTree, displays: dict[str, _DataTreeDisplay]
+) -> list[str]:
+    display = displays[node.path]
+    sections = []
+    if node.children and not display.disabled:
+        sections.append(children_section(node.children, displays))
+    sections.extend(display.sections)
+    return sections
 
-def datatree_child_repr(node: DataTree, end: bool = False) -> str:
+
+def datatree_child_repr(
+    node: DataTree,
+    displays: dict[str, _DataTreeDisplay],
+    end: bool,
+) -> str:
     # Wrap DataTree HTML representation with a tee to the left of it.
     #
     # Enclosing HTML tag is a <div> with :code:`display: inline-grid` style.
@@ -487,39 +576,43 @@ def datatree_child_repr(node: DataTree, end: bool = False) -> str:
     # └─ [    title    ]
     #    |   details   |
     #    |_____________|
-    end = bool(end)
-    height = "100%" if end is False else "1.2em"  # height of line
+
+    vline_height = "1.2em" if end else "100%"
 
     path = escape(node.path)
-    sections = datatree_node_sections(node, root=False)
-    section_items = "".join(f"<li class='xr-section-item'>{s}</li>" for s in sections)
+    display = displays[node.path]
+
+    group_id = "group-" + str(uuid.uuid4())
+    collapsed = " checked" if display.collapsed else ""
+    tip = " title='Expand/collapse group'" if not display.disabled else ""
+
+    sections = datatree_sections(node, displays)
+    sections_html = _sections_repr(sections) if sections else ""
 
-    # TODO: Can we make the group name clickable to toggle the sections below?
-    # This looks like it would require the input/label pattern used above.
     html = f"""
         <div class='xr-group-box'>
-            <div class='xr-group-box-vline' style='height: {height}'></div>
+            <div class='xr-group-box-vline' style='height: {vline_height}'></div>
             <div class='xr-group-box-hline'></div>
             <div class='xr-group-box-contents'>
-                <div class='xr-header'>
-                    <div class='xr-group-name'>{path}</div>
-                </div>
-                <ul class='xr-sections'>
-                    {section_items}
-                </ul>
+                <input id='{group_id}' type='checkbox'{collapsed} />
+                <label for='{group_id}'{tip}>
+                    {path}
+                    <span>({display.item_count})</span>
+                </label>
+                {sections_html}
             </div>
         </div>
     """
     return "".join(t.strip() for t in html.split("\n"))
 
 
 def datatree_repr(node: DataTree) -> str:
+    displays = _build_datatree_displays(node)
     header_components = [
         f"<div class='xr-obj-type'>xarray.{type(node).__name__}</div>",
     ]
     if node.name is not None:
         name = escape(repr(node.name))
         header_components.append(f"<div class='xr-obj-name'>{name}</div>")
-
-    sections = datatree_node_sections(node, root=True)
+    sections = datatree_sections(node, displays)
     return _obj_repr(node, header_components, sections)