Use ipytree in .tree representation (#450)

Author: jakirkham

docs/tutorial.rst

@@ -424,7 +424,8 @@ Groups also have the :func:`zarr.hierarchy.Group.tree` method, e.g.::
          ├── bar (1000000,) int64
          └── baz (1000, 1000) float32
 
-If you're using Zarr within a Jupyter notebook, calling ``tree()`` will generate an
+If you're using Zarr within a Jupyter notebook (requires
+`ipytree <https://github.com/QuantStack/ipytree>`_), calling ``tree()`` will generate an
 interactive tree representation, see the `repr_tree.ipynb notebook
 <http://nbviewer.jupyter.org/github/zarr-developers/zarr-python/blob/master/notebooks/repr_tree.ipynb>`_
 for more examples.

notebooks/repr_tree.ipynb

@@ -1,6 +1,8 @@
 # optional library requirements
 bsddb3==6.2.6; sys_platform != 'win32'
 lmdb==0.97; sys_platform != 'win32'
+# optional library requirements for Jupyter
+ipytree==0.1.3
 # optional library requirements for services
 # don't let pyup change pinning for azure-storage-blob, need to pin to older
 # version to get compatibility with azure storage emulator on appveyor

requirements_dev_optional.txt

@@ -30,6 +30,11 @@
         'setuptools>=38.6.0',
         'setuptools-scm>1.5.4'
     ],
+    extras_require={
+        'jupyter': [
+            'ipytree',
+        ],
+    },
     python_requires='>=3.5',
     install_requires=dependencies,
     package_dir={'': '.'},

setup.py

@@ -10,6 +10,12 @@
 
 import numpy as np
 import pytest
+
+try:
+    import ipytree
+except ImportError:  # pragma: no cover
+    ipytree = None
+
 from numcodecs import Zlib
 from numpy.testing import assert_array_equal
 
@@ -1199,11 +1205,10 @@ def _check_tree(g, expect_bytes, expect_text):
     assert expect_text == str(g.tree())
     expect_repr = expect_text
     assert expect_repr == repr(g.tree())
-    # test _repr_html_ lightly
-    # noinspection PyProtectedMember
-    html = g.tree()._repr_html_().strip()
-    assert html.startswith('<link')
-    assert html.endswith('</script>')
+    if ipytree:
+        # noinspection PyProtectedMember
+        widget = g.tree()._ipython_display_()
+        isinstance(widget, ipytree.Tree)
 
 
 def test_tree():

zarr/tests/test_hierarchy.py

@@ -5,7 +5,8 @@
 from zarr.util import (guess_chunks, human_readable_size, info_html_report,
                        info_text_report, is_total_slice, normalize_chunks,
                        normalize_fill_value, normalize_order,
-                       normalize_resize_args, normalize_shape)
+                       normalize_resize_args, normalize_shape,
+                       tree_array_icon, tree_group_icon, tree_get_icon)
 
 
 def test_normalize_shape():
@@ -152,3 +153,10 @@ def test_info_html_report():
     actual = info_html_report(items)
     assert '<table' == actual[:6]
     assert '</table>' == actual[-8:]
+
+
+def test_tree_get_icon():
+    assert tree_get_icon("Array") == tree_array_icon
+    assert tree_get_icon("Group") == tree_group_icon
+    with pytest.raises(ValueError):
+        tree_get_icon("Baz")

zarr/tests/test_util.py

@@ -3,8 +3,7 @@
 import json
 import math
 import numbers
-import uuid
-from textwrap import TextWrapper, dedent
+from textwrap import TextWrapper
 
 import numpy as np
 from asciitree import BoxStyle, LeftAligned
@@ -399,83 +398,43 @@ def get_text(self, node):
         return node.get_text()
 
 
-def tree_html_sublist(node, root=False, expand=False):
-    result = ''
-    data_jstree = '{"type": "%s"}' % node.get_type()
-    if root or (expand is True) or (isinstance(expand, int) and node.depth < expand):
-        css_class = 'jstree-open'
-    else:
-        css_class = ''
-    result += "<li data-jstree='{}' class='{}'>".format(data_jstree, css_class)
-    result += '<span>{}</span>'.format(node.get_text())
-    children = node.get_children()
-    if children:
-        result += '<ul>'
-        for c in children:
-            result += tree_html_sublist(c, expand=expand)
-        result += '</ul>'
-    result += '</li>'
-    return result
+tree_group_icon = 'folder'
+tree_array_icon = 'table'
 
 
-def tree_html(group, expand, level):
+def tree_get_icon(stype):
+    if stype == "Array":
+        return tree_array_icon
+    elif stype == "Group":
+        return tree_group_icon
+    else:
+        raise ValueError("Unknown type: %s" % stype)
 
-    result = ''
 
-    # include CSS for jstree default theme
-    css_url = '//cdnjs.cloudflare.com/ajax/libs/jstree/3.3.3/themes/default/style.min.css'
-    result += '<link rel="stylesheet" href="{}"/>'.format(css_url)
+def tree_widget_sublist(node, root=False, expand=False):
+    import ipytree
 
-    # construct the tree as HTML nested lists
-    node_id = uuid.uuid4()
-    result += '<div id="{}" class="zarr-tree">'.format(node_id)
-    result += '<ul>'
-    root = TreeNode(group, level=level)
-    result += tree_html_sublist(root, root=True, expand=expand)
-    result += '</ul>'
-    result += '</div>'
-
-    # construct javascript
-    result += dedent("""
-        <script>
-            if (!require.defined('jquery')) {
-                require.config({
-                    paths: {
-                        jquery: '//cdnjs.cloudflare.com/ajax/libs/jquery/1.12.1/jquery.min'
-                    },
-                });
-            }
-            if (!require.defined('jstree')) {
-                require.config({
-                    paths: {
-                        jstree: '//cdnjs.cloudflare.com/ajax/libs/jstree/3.3.3/jstree.min'
-                    },
-                });
-            }
-            require(['jstree'], function() {
-                $('#%s').jstree({
-                    types: {
-                        Group: {
-                            icon: "%s"
-                        },
-                        Array: {
-                            icon: "%s"
-                        }
-                    },
-                    plugins: ["types"]
-                });
-            });
-        </script>
-    """ % (node_id, tree_group_icon, tree_array_icon))
+    result = ipytree.Node()
+    result.icon = tree_get_icon(node.get_type())
+    if root or (expand is True) or (isinstance(expand, int) and node.depth < expand):
+        result.opened = True
+    else:
+        result.opened = False
+    result.name = node.get_text()
+    result.nodes = [tree_widget_sublist(c, expand=expand) for c in node.get_children()]
+    result.disabled = True
 
     return result
 
 
-tree_group_icon = 'fa fa-folder'
-tree_array_icon = 'fa fa-table'
-# alternatives...
-# tree_group_icon: 'jstree-folder'
-# tree_array_icon: 'jstree-file'
+def tree_widget(group, expand, level):
+    import ipytree
+
+    result = ipytree.Tree()
+    root = TreeNode(group, level=level)
+    result.add_node(tree_widget_sublist(root, root=True, expand=expand))
+
+    return result
 
 
 class TreeViewer(object):
@@ -531,8 +490,10 @@ def __unicode__(self):
     def __repr__(self):
         return self.__unicode__()
 
-    def _repr_html_(self):
-        return tree_html(self.group, expand=self.expand, level=self.level)
+    def _ipython_display_(self):
+        tree = tree_widget(self.group, expand=self.expand, level=self.level)
+        tree._ipython_display_()
+        return tree
 
 
 def check_array_shape(param, array, shape):