"
+
+
def _obj_repr(obj, header_components, sections):
"""Return HTML repr of an xarray object.
@@ -301,7 +314,6 @@ def _obj_repr(obj, header_components, sections):
"""
header = f"
{''.join(h for h in header_components)}
"
- sections = "".join(f"
{s}
" for s in sections)
icons_svg, css_style = _load_static_files()
return (
@@ -310,7 +322,7 @@ def _obj_repr(obj, header_components, sections):
f"
{escape(repr(obj))}
"
"
"
f"{header}"
- f"
{sections}
"
+ f"{_sections_repr(sections)}"
"
"
""
)
@@ -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("
...
")
-
- return "".join(
- [
- "
",
- "".join(children_html),
- "
",
- ]
+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(
+ "Too many items to display (display_max_html_elements exceeded)",
+ 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"
{children_html}
"
+
+
+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
with :code:`display: inline-grid` style.
@@ -487,26 +576,30 @@ 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"
{s}
" 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"""
" in result
- else:
- assert f"file_{i}" not in result
-
- assert "6/25" in result
- for i in range(number_of_groups):
- if i < 3 or i >= (number_of_groups - 3):
- assert f"group_{i}" in result
- else:
- assert f"group_{i}" not in result
-
- with xr.set_options(display_style="html", display_max_children=3):
- result = tree._repr_html_()
-
- assert "3/20" in result
- for i in range(number_of_files):
- if i < 2 or i >= (number_of_files - 1):
- assert f"file_{i}" in result
- else:
- assert f"file_{i}" not in result
-
- assert "3/25" in result
- for i in range(number_of_groups):
- if i < 2 or i >= (number_of_groups - 1):
- assert f"group_{i}" in result
- else:
- assert f"group_{i}" not in result
+
+ with xr.set_options(display_max_html_elements=25):
+ result = xarray_html_only_repr(tree)
+ assert result.count("file_0/group_9") == 1
+ assert result.count("file_1/group_0") == 0 # disabled
+ assert result.count("Too many items to display") == 9 + 10
+
+ with xr.set_options(display_max_html_elements=1000):
+ result = xarray_html_only_repr(tree)
+ assert result.count("Too many items to display") == 0
class TestDataTreeInheritance: