html repr (#78)

* html repr displays data in root group

* displays under name 'xarray.DataTree'

* creates a html repr for each sub-group, but it looks messed up

* correctly indents sub-groups

* show names of groups

* refactoring

* dodge type hinting bug

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix failing test in merg

* fix bug caused by merge

* whatsnew

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: TomNicholas

datatree/datatree.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections import OrderedDict
+from html import escape
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -16,9 +17,10 @@
 
 from xarray import DataArray, Dataset
 from xarray.core import utils
+from xarray.core.options import OPTIONS as XR_OPTS
 from xarray.core.variable import Variable
 
-from .formatting import tree_repr
+from . import formatting, formatting_html
 from .mapping import TreeIsomorphismError, check_isomorphic, map_over_subtree
 from .ops import (
     DataTreeArithmeticMixin,
@@ -189,11 +191,17 @@ def _pre_attach(self: DataTree, parent: DataTree) -> None:
                 f"parent {parent.name} already contains a data variable named {self.name}"
             )
 
-    def __repr__(self):
-        return tree_repr(self)
+    def __repr__(self) -> str:
+        return formatting.datatree_repr(self)
 
-    def __str__(self):
-        return tree_repr(self)
+    def __str__(self) -> str:
+        return formatting.datatree_repr(self)
+
+    def _repr_html_(self):
+        """Make html representation of datatree object"""
+        if XR_OPTS["display_style"] == "text":
+            return f"<pre>{escape(repr(self))}</pre>"
+        return formatting_html.datatree_repr(self)
 
     def get(
         self: DataTree, key: str, default: Optional[DataTree | DataArray] = None
@@ -227,8 +235,10 @@ def __getitem__(self: DataTree, key: str) -> DataTree | DataArray:
         key : str
             Name of variable / node, or unix-like path to variable / node.
         """
+
         # Either:
         if utils.is_dict_like(key):
+
             # dict-like indexing
             raise NotImplementedError("Should this index over whole tree?")
         elif isinstance(key, str):
@@ -243,7 +253,7 @@ def __getitem__(self: DataTree, key: str) -> DataTree | DataArray:
                 "implemented via .subset"
             )
         else:
-            raise ValueError("Invalid format for key")
+            raise ValueError(f"Invalid format for key: {key}")
 
     def _set(self, key: str, val: DataTree | CoercibleValue) -> None:
         """
@@ -352,6 +362,13 @@ def from_dict(
     def nbytes(self) -> int:
         return sum(node.ds.nbytes if node.has_data else 0 for node in self.subtree)
 
+    def __len__(self) -> int:
+        if self.children:
+            n_children = len(self.children)
+        else:
+            n_children = 0
+        return n_children + len(self.ds)
+
     def isomorphic(
         self,
         other: DataTree,

datatree/formatting.py

@@ -52,7 +52,7 @@ def diff_tree_repr(a, b, compat):
     return "\n".join(summary)
 
 
-def tree_repr(dt):
+def datatree_repr(dt):
     """A printable representation of the structure of this entire tree."""
     renderer = RenderTree(dt)
 

datatree/formatting_html.py

@@ -0,0 +1,57 @@
+from functools import partial
+from html import escape
+from typing import Any, Mapping
+
+from xarray.core.formatting_html import (
+    _mapping_section,
+    _obj_repr,
+    attr_section,
+    coord_section,
+    datavar_section,
+    dim_section,
+)
+from xarray.core.options import OPTIONS
+
+OPTIONS["display_expand_groups"] = "default"
+
+
+def summarize_children(children: Mapping[str, Any]) -> str:
+    children_li = "".join(
+        f"<ul class='xr-sections'>{node_repr(n, c)}</ul>" for n, c in children.items()
+    )
+
+    return (
+        "<ul class='xr-sections'>"
+        f"<div style='padding-left:2rem;'>{children_li}<br></div>"
+        "</ul>"
+    )
+
+
+children_section = partial(
+    _mapping_section,
+    name="Groups",
+    details_func=summarize_children,
+    max_items_collapse=1,
+    expand_option_name="display_expand_groups",
+)
+
+
+def node_repr(group_title: str, dt: Any) -> str:
+    header_components = [f"<div class='xr-obj-type'>{escape(group_title)}</div>"]
+
+    ds = dt.ds
+
+    sections = [
+        children_section(dt.children),
+        dim_section(ds),
+        coord_section(ds.coords),
+        datavar_section(ds.data_vars),
+        attr_section(ds.attrs),
+    ]
+
+    return _obj_repr(ds, header_components, sections)
+
+
+def datatree_repr(dt: Any) -> str:
+    obj_type = f"datatree.{type(dt).__name__}"
+    return node_repr(obj_type, dt)

datatree/treenode.py

@@ -134,7 +134,7 @@ def _detach(self, parent: Tree | None) -> None:
     def _attach(self, parent: Tree | None, child_name: str = None) -> None:
         if parent is not None:
             if child_name is None:
-                raise ValueError()
+                raise ValueError("Cannot directly assign a parent to an unnamed node")
 
             self._pre_attach(parent)
             parentchildren = parent._children

docs/source/whats-new.rst

@@ -32,7 +32,9 @@ New Features
   root node via ``"/"``, and the current node via ``"."``. (Internally it actually uses ``pathlib`` now.)
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
 - New path-like API methods, such as ``.relative_to``, ``.find_common_ancestor``, and ``.same_tree``.
-- Some new diction-like methods, such as ``DataTree.get`` and ``DataTree.update``. (:pull:`76`)
+- Some new dictionary-like methods, such as ``DataTree.get`` and ``DataTree.update``. (:pull:`76`)
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
+- New HTML repr, which will automatically display in a jupyter notebook. (:pull:`78`)
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
 
 Breaking changes