Add better ipytree error handling in jupyter notebooks (#539)

* Add better ipytree error handling in jupyter notebooks

* Fix pep8 issues

* Use notebook instead of jupyter in setup.py

* Update error with jupyterlab specific fix

* Update error message in test function too

* Update docs/release.rst

* Update error message with conda

Author: mzjp2

docs/release.rst

@@ -7,6 +7,9 @@ Next release
 
 * Fix minor bug in `N5Store`. 
   By :user:`gsakkis`, :issue:`550`.
+* Improve error message in Jupyter when trying to use the ``ipytree`` widget
+  without ``ipytree`` installed.
+  By :user:`Zain Patel <mzjp2>; :issue:`537`
 
 
 .. _release_2.4.0:

setup.py

@@ -24,14 +24,15 @@
     use_scm_version={
         'version_scheme': 'guess-next-dev',
         'local_scheme': 'dirty-tag',
-        'write_to': 'zarr/version.py'
+        'write_to': 'zarr/version.py',
     },
     setup_requires=[
         'setuptools>=38.6.0',
-        'setuptools-scm>1.5.4'
+        'setuptools-scm>1.5.4',
     ],
     extras_require={
         'jupyter': [
+            'notebook',
             'ipytree',
         ],
     },

zarr/tests/test_util.py

@@ -1,12 +1,16 @@
 # -*- coding: utf-8 -*-
+import re
+from unittest import mock
+
 import numpy as np
 import pytest
 
 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,
-                       tree_array_icon, tree_group_icon, tree_get_icon)
+                       tree_array_icon, tree_group_icon, tree_get_icon,
+                       tree_widget)
 
 
 def test_normalize_shape():
@@ -160,3 +164,15 @@ def test_tree_get_icon():
     assert tree_get_icon("Group") == tree_group_icon
     with pytest.raises(ValueError):
         tree_get_icon("Baz")
+
+
+@mock.patch.dict("sys.modules", {"ipytree": None})
+def test_tree_widget_missing_ipytree():
+    pattern = (
+        "Run `pip install zarr[jupyter]` or `conda install ipytree`"
+        "to get the required ipytree dependency for displaying the tree "
+        "widget. If using jupyterlab, you also need to run "
+        "`jupyter labextension install ipytree`"
+        )
+    with pytest.raises(ImportError, match=re.escape(pattern)):
+        tree_widget(None, None, None)

zarr/util.py

@@ -428,7 +428,15 @@ def tree_widget_sublist(node, root=False, expand=False):
 
 
 def tree_widget(group, expand, level):
-    import ipytree
+    try:
+        import ipytree
+    except ImportError as error:
+        raise ImportError(
+            "{}: Run `pip install zarr[jupyter]` or `conda install ipytree`"
+            "to get the required ipytree dependency for displaying the tree "
+            "widget. If using jupyterlab, you also need to run "
+            "`jupyter labextension install ipytree`".format(error)
+        )
 
     result = ipytree.Tree()
     root = TreeNode(group, level=level)