feat: add rich traceback handler as default (#215)

hanxiao · web-flow · commit 1ae1ffa8179b · 2022-03-19T15:03:10.000+01:00
diff --git a/docarray/__init__.py b/docarray/__init__.py
@@ -1,4 +1,11 @@
 __version__ = '0.9.15'
 
+import os
+
 from .document import Document
 from .array import DocumentArray
+
+if 'DA_NO_RICH_HANDLER' not in os.environ:
+    from rich.traceback import install
+
+    install()
diff --git a/docarray/array/mixins/io/pushpull.py b/docarray/array/mixins/io/pushpull.py
@@ -127,7 +127,12 @@ def pull(
         url = f'{_get_cloud_api()}/v2/rpc/da.pull?token={token}'
         response = requests.get(url)
 
-        url = response.json()['data']['download']
+        if response.ok:
+            url = response.json()['data']['download']
+        else:
+            raise FileNotFoundError(
+                f'can not find `{token}` DocumentArray on the Cloud. Misspelled?'
+            )
 
         with requests.get(
             url,
@@ -161,26 +166,3 @@ def pull(
                     fp.write(_source.content)
 
             return r
-
-
-def _get_progressbar(show_progress):
-    from rich.progress import (
-        BarColumn,
-        DownloadColumn,
-        Progress,
-        TimeRemainingColumn,
-        TransferSpeedColumn,
-    )
-
-    return Progress(
-        BarColumn(),
-        "[progress.percentage]{task.percentage:>3.1f}%",
-        "&bull;",
-        DownloadColumn(),
-        "&bull;",
-        TransferSpeedColumn(),
-        "&bull;",
-        TimeRemainingColumn(),
-        transient=True,
-        disable=not show_progress,
-    )
diff --git a/docarray/document/mixins/plot.py b/docarray/document/mixins/plot.py
@@ -8,26 +8,40 @@ def _ipython_display_(self):
         """Displays the object in IPython as a side effect"""
         self.summary()
 
+    def __rich_console__(self, console, options):
+        yield f":page_facing_up: [b]Document[/b]: [cyan]{self.id}[cyan]"
+        from rich.table import Table
+
+        my_table = Table('Attribute', 'Value')
+        for f in self.non_empty_fields:
+            if f not in ('id', 'chunks', 'matches'):
+                my_table.add_row(f, str(getattr(self, f)))
+        if my_table.rows:
+            yield my_table
+
     def summary(self) -> None:
         """ Print non-empty fields and nested structure of this Document object."""
-        _str_list = []
-        self._plot_recursion(_str_list, indent=0)
-        print('\n'.join(_str_list))
+        from rich import print
 
-    def _plot_recursion(self, _str_list, indent, box_char='&#9500;&#9472;'):
-        prefix = (' ' * indent + box_char) if indent else ''
-        _str_list.append(f'{prefix} {self}')
+        print(self._plot_recursion())
 
+    def _plot_recursion(self, tree=None):
+        if tree is None:
+            from rich.tree import Tree
+
+            tree = Tree(self)
+        else:
+            tree = tree.add(self)
         for a in ('matches', 'chunks'):
             if getattr(self, a):
-                prefix = ' ' * (indent + 4) + '&#9492;&#9472;'
-                _str_list.append(f'{prefix} {a}')
-
-                for d in getattr(self, a)[:-1]:
-                    d._plot_recursion(_str_list, indent=len(prefix) + 4)
-                getattr(self, a)[-1]._plot_recursion(
-                    _str_list, indent=len(prefix) + 4, box_char='&#9492;&#9472;'
-                )
+                if a == 'chunks':
+                    _icon = ':diamond_with_a_dot:'
+                else:
+                    _icon = ':large_orange_diamond:'
+                _match_tree = tree.add(f'{_icon} [b]{a.capitalize()}[/b]')
+                for d in getattr(self, a):
+                    d._plot_recursion(_match_tree)
+        return tree
 
     def display(self):
         """ Plot image data from :attr:`.tensor` or :attr:`.uri`. """
diff --git a/tests/unit/array/mixins/test_pushpull.py b/tests/unit/array/mixins/test_pushpull.py
@@ -18,6 +18,7 @@ class PullMockResponse:
     def __init__(self, status_code: int = 200):
         self.status_code = status_code
         self.headers = {'Content-length': 1}
+        self.ok = True
 
     def json(self):
         return {
