feat(array): auto render in colab and jupyter (#40)

hanxiao · web-flow · commit 699c041e03bb · 2022-01-12T15:02:42.000+01:00
diff --git a/docarray/array/mixins/io/csv.py b/docarray/array/mixins/io/csv.py
@@ -26,7 +26,8 @@ def save_embeddings_csv(self, file: Union[str, TextIO], **kwargs) -> None:
             file_ctx = nullcontext(file)
         else:
             file_ctx = open(file, 'w')
-        np.savetxt(file_ctx, self.embeddings, **kwargs)
+        with file_ctx:
+            np.savetxt(file_ctx, self.embeddings, **kwargs)
 
     def save_csv(
         self,
diff --git a/docarray/array/mixins/plot.py b/docarray/array/mixins/plot.py
@@ -3,6 +3,7 @@
 import os.path
 import tempfile
 import threading
+import time
 import warnings
 from collections import Counter
 from math import sqrt, ceil, floor
@@ -220,21 +221,57 @@ def _get_fastapi_app():
                 kwargs=dict(app=app, port=port, log_level='error'),
                 daemon=True,
             )
-            t_m.start()
             url_html_path = (
                 f'http://localhost:{port}/static/index.html?config={config_fn}'
             )
+            t_m.start()
             try:
-                import webbrowser
-
-                webbrowser.open(url_html_path, new=2)
+                _env = str(get_ipython())  # noqa
+                if 'ZMQInteractiveShell' in _env:
+                    _env = 'jupyter'
+                elif 'google.colab' in _env:
+                    _env = 'colab'
             except:
-                pass  # intentional pass, browser support isn't cross-platform
-            finally:
-                print(
-                    f'You should see a webpage opened in your browser, '
-                    f'if not, you may open {url_html_path} manually'
+                _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))
+            elif _env == 'colab':
+                from google.colab.output import eval_js  # noqa
+
+                colab_url = eval_js(f'google.colab.kernel.proxyPort({port})')
+                colab_url += f'/static/index.html?config={config_fn}'
+                warnings.warn(
+                    f'Showing iframe in cell, you may want to open {colab_url} in a new tab for better experience. '
+                    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
+
+                display(IFrame(src=colab_url, width="100%", height=600))
+            elif _env == 'local':
+                try:
+                    import webbrowser
+
+                    webbrowser.open(url_html_path, new=2)
+                except:
+                    pass  # intentional pass, browser support isn't cross-platform
+                finally:
+                    print(
+                        f'You should see a webpage opened in your browser, '
+                        f'if not, you may open {url_html_path} manually'
+                    )
             t_m.join()
         return path
 
diff --git a/setup.py b/setup.py
@@ -52,6 +52,7 @@
             'scipy',
             'av',
             'fastapi',
+            'uvicorn',
         ]
     },
     classifiers=[

Original file line number	Diff line number	Diff line change
`@@ -52,6 +52,7 @@`
`52`	`52`	`'scipy',`
`53`	`53`	`'av',`
`54`	`54`	`'fastapi',`
	`55`	`+ 'uvicorn',`
`55`	`56`	`]`
`56`	`57`	`},`
`57`	`58`	`classifiers=[`