feat: better jupyter notebook support (#44)

docarray · Jan 13, 2022 · edf836d · edf836d
1 parent 995dbd3
commit edf836d
Show file tree

Hide file tree

Showing 17 changed files with 143 additions and 23 deletions.
diff --git a/docarray/array/mixins/plot.py b/docarray/array/mixins/plot.py
@@ -11,10 +11,16 @@
 
 import numpy as np
 
+from ...helper import deprecate_by
+
 
 class PlotMixin:
     """Helper functions for plotting the arrays. """
 
+    def _ipython_display_(self):
+        """Displays the object in IPython as a side effect"""
+        self.summary()
+
     def summary(self):
         """Print the structure and attribute summary of this DocumentArray object.
 
@@ -234,17 +240,17 @@ def _get_fastapi_app():
             except:
                 _env = 'local'
             if _env == 'jupyter':
-                warnings.warn(
-                    f'Showing iframe in cell, you may want to open {url_html_path} in a new tab for better experience. '
-                    f'Also, `localhost` may need to be changed to the IP address if your jupyter is running remotely. '
-                    f'Click "stop" button in the toolbar to move to the next cell.'
-                )
                 time.sleep(
                     1
                 )  # jitter is required otherwise encouter werid `strict-origin-when-cross-origin` error in browser
                 from IPython.display import IFrame, display  # noqa
 
                 display(IFrame(src=url_html_path, width="100%", height=600))
+                warnings.warn(
+                    f'Showing iframe in cell, you may want to open {url_html_path} in a new tab for better experience. '
+                    f'Also, `localhost` may need to be changed to the IP address if your jupyter is running remotely. '
+                    f'Click "stop" button in the toolbar to move to the next cell.'
+                )
             elif _env == 'colab':
                 from google.colab.output import eval_js  # noqa
 

diff --git a/docarray/document/mixins/plot.py b/docarray/document/mixins/plot.py
@@ -1,7 +1,4 @@
-import base64
-from typing import Optional
-
-from ...helper import random_identity, download_mermaid_url
+from ...helper import deprecate_by
 
 
 class PlotMixin:
@@ -32,15 +29,58 @@ def _plot_recursion(self, _str_list, indent, box_char='├─'):
                     _str_list, indent=len(prefix) + 4, box_char='└─'
                 )
 
-    def plot(self):
+    def display(self):
         """ Plot image data from :attr:`.blob` or :attr:`.uri`. """
         from IPython.display import Image, display
 
-        if self.blob is not None:
-            import PIL.Image
+        if self.uri:
+            if self.mime_type.startswith('audio'):
+                _html5_audio_player(self.uri)
+            elif self.mime_type.startswith('video'):
+                _html5_video_player(self.uri)
+            else:
+                display(Image(self.uri))
+        elif self.blob is not None:
+            try:
+                import PIL.Image
+
+                p = PIL.Image.fromarray(self.blob)
+                if p.mode != 'RGB':
+                    raise
+                display(p)
+            except:
+                import matplotlib.pyplot as plt
 
-            display(PIL.Image.fromarray(self.blob))
-        elif self.uri:
-            display(Image(self.uri))
+                plt.matshow(self.blob)
         else:
-            raise ValueError('`uri` and `blob` is empty')
+            self.summary()
+
+    plot = deprecate_by(display, removed_at='0.5')
+
+
+def _html5_video_player(uri):
+    from IPython.core.display import HTML  # noqa
+
+    src = f'''
+    <body>
+    <video width="320" height="240" autoplay muted controls>
+    <source src="files/{uri}">
+    Your browser does not support the video tag.
+    </video>
+    </body>
+    '''
+    display(HTML(src))  # noqa
+
+
+def _html5_audio_player(uri):
+    from IPython.core.display import HTML  # noqa
+
+    src = f'''
+    <body>
+    <audio controls="controls" style="width:320px" >
+      <source src="files/{uri}"/>
+      Your browser does not support the audio element.
+    </audio>
+    </body>
+    '''
+    display(HTML(src))  # noqa
diff --git a/docarray/helper.py b/docarray/helper.py
@@ -30,12 +30,13 @@ def typename(obj):
         return str(obj)
 
 
-def deprecate_by(new_fn):
+def deprecate_by(new_fn, removed_at: str):
     """A helper function to label deprecated function
 
     Usage: old_fn_name = deprecate_by(new_fn)
 
     :param new_fn: the new function
+    :param removed_at: removed at which version
     :return: a wrapped function with old function name
     """
 
@@ -44,9 +45,8 @@ def _f(*args, **kwargs):
 
         old_fn_name = inspect.stack()[1][4][0].strip().split("=")[0].strip()
         warnings.warn(
-            f'`{old_fn_name}` is renamed to `{new_fn.__name__}` with the same usage, please use the latter instead. '
-            f'The old function will be removed soon.',
-            DeprecationWarning,
+            f'`{old_fn_name}` is renamed to `.{new_fn.__name__}()` with the same usage, please use the latter instead. The old function will be removed in {removed_at}.',
+            FutureWarning,
         )
         return new_fn(*args, **kwargs)
 
@@ -294,15 +294,19 @@ def decompress_bytes(data: bytes, algorithm: Optional[str] = None) -> bytes:
 def get_compress_ctx(algorithm: Optional[str] = None, mode: str = 'wb'):
     if algorithm == 'lz4':
         import lz4.frame
+
         compress_ctx = lambda x: lz4.frame.LZ4FrameFile(x, mode)
     elif algorithm == 'gzip':
         import gzip
+
         compress_ctx = lambda x: gzip.GzipFile(fileobj=x, mode=mode)
     elif algorithm == 'bz2':
         import bz2
+
         compress_ctx = lambda x: bz2.BZ2File(x, mode)
     elif algorithm == 'lzma':
         import lzma
+
         compress_ctx = lambda x: lzma.LZMAFile(x, mode)
     else:
         compress_ctx = None

diff --git a/docs/fundamentals/document/visualization.md b/docs/fundamentals/document/visualization.md
@@ -1,6 +1,6 @@
 # Visualization
 
-If you have an image Document (with possible image data in `.uri`/`.blob`), you can directly visualize it via {meth}`~docarray.document.mixins.plot.PlotMixin.plot`.
+If you have an image Document (with possible image data in `.uri`/`.blob`), you can directly visualize it via {meth}`~docarray.document.mixins.plot.PlotMixin.display`.
 
 ```{figure} images/doc-plot-in-jupyter.png
 ```

diff --git a/docs/fundamentals/notebook-support/audio-video.png b/docs/fundamentals/notebook-support/audio-video.png
diff --git a/docs/fundamentals/notebook-support/doc-array.png b/docs/fundamentals/notebook-support/doc-array.png
diff --git a/docs/fundamentals/notebook-support/embedding-ani1.gif b/docs/fundamentals/notebook-support/embedding-ani1.gif
diff --git a/docs/fundamentals/notebook-support/embedding-ani2.gif b/docs/fundamentals/notebook-support/embedding-ani2.gif
diff --git a/docs/fundamentals/notebook-support/image-blob.png b/docs/fundamentals/notebook-support/image-blob.png
diff --git a/docs/fundamentals/notebook-support/image-doc-blob.png b/docs/fundamentals/notebook-support/image-doc-blob.png
diff --git a/docs/fundamentals/notebook-support/image-doc.png b/docs/fundamentals/notebook-support/image-doc.png
diff --git a/docs/fundamentals/notebook-support/image-sprite.png b/docs/fundamentals/notebook-support/image-sprite.png
diff --git a/docs/fundamentals/notebook-support/index.md b/docs/fundamentals/notebook-support/index.md
@@ -0,0 +1,69 @@
+# Notebook/Colab Support
+
+Many data scientists work with Jupyter Notebook or Google Colab especially at the early prototyping stage. We understand that and that's why we optimize DocArray's user experience in these two environments. In this chapter, I use Jupyter Notebook as an example to demonstrate the features that can improve your productivity. On Google Colab, it is the same experience.
+
+```{tip}
+Some features require extra dependencies beyond the basic install of DocArray. Use `pip install "docarray[full]"` to enable them.
+```
+
+## Display Document
+
+A cell with a Document object will be pretty-printed with its non-empty field and `id`.
+
+```{figure} single-doc.png
+```
+
+If a Document is nested, then it pretty-prints the nested structure.
+
+```{figure} single-doc-nested.png
+```
+
+### Display rich content
+
+If a Document is an image Document, you can use {meth}`~docarray.document.mixins.plot.PlotMixin.display` to visualize it.
+
+```{figure} image-doc.png
+```
+
+Note that it finds `.blob` or `.uri` for visualization.
+
+```{figure} image-doc-blob.png
+```
+
+This works even if your Document is not a real image but just a `ndarray` in `.blob`.
+
+```{figure} image-blob.png
+```
+
+Video and audio Document can be displayed as well, you can play them in the cell.
+
+```{figure} audio-video.png
+```
+
+## Display DocumentArray
+
+A cell with a DocumentArray object can be pretty-printed automatically.
+
+```{figure} doc-array.png
+```
+
+### Display image sprite
+
+DocumentArray with all image Documents (image is either in `.uri` or `.blob`) can be plotted into one sprite image.
+
+```{figure} image-sprite.png
+```
+
+### Display embeddings
+
+DocumentArray with non-empty `.embeddings` can be visualized interactively via {meth}`~docarray.array.mixins.plot.PlotMixin.plot_embeddings`
+
+```{figure} embedding-ani1.gif
+```
+
+
+DocumentArray with non-empty `.embeddings`  and image Documents can be visualized in a much richer way via `.plot_embeddings(image_sprites=True)`.
+
+```{figure} embedding-ani2.gif
+```
+
diff --git a/docs/fundamentals/notebook-support/single-doc-nested.png b/docs/fundamentals/notebook-support/single-doc-nested.png
diff --git a/docs/fundamentals/notebook-support/single-doc.png b/docs/fundamentals/notebook-support/single-doc.png
diff --git a/docs/index.md b/docs/index.md
@@ -92,6 +92,7 @@ get-started/what-is
 
 fundamentals/document/index
 fundamentals/documentarray/index
+fundamentals/notebook-support/index
 datatypes/index
 ```
 

diff --git a/tests/unit/document/test_summary.py b/tests/unit/document/test_summary.py
@@ -21,9 +21,9 @@ def test_single_doc_summary():
 
 def test_plot_image():
     d = Document(uri=os.path.join(cur_dir, 'toydata/test.png'))
-    d.plot()
+    d.display()
 
     d.load_uri_to_image_blob()
     d.uri = None
 
-    d.plot()
+    d.display()