subclass IPython.core.display HTML object for more flexibility.

plotly/tests/test_optional/test_ipython/test_embed.py

@@ -0,0 +1,54 @@
+from __future__ import absolute_import
+
+import plotly.tools as tls
+import imghdr
+import threading
+import unittest
+
+
+class TestPlotlyDisplay(unittest.TestCase):
+
+    def setUp(self):
+        plot_info = {"un": "plotlyimagetest", "fid": "2"}
+        url = "https://plot.ly/~{un}/{fid}".format(**plot_info)
+        self.display_obj = tls.embed(url)
+        self.results = {}
+        self.threads = []
+        self.format_to_func = {
+            "jpeg": self.jpeg_worker,
+            "png": self.png_worker,
+            "svg": self.svg_worker,
+            "pdf": self.pdf_worker}
+
+    def test_plotly_display(self):
+        for f_format, func in self.format_to_func.items():
+            self.threads += [threading.Thread(target=func)]
+            self.threads[-1].setDaemon(True)
+            self.threads[-1].start()
+        for thread in self.threads:
+            thread.join()
+        for f_format in self.format_to_func:
+            result = self.results.get(f_format, False)
+            print("{f_format}: {result}".format(f_format=f_format,
+                                                result=result))
+            assert self.results.get(f_format)
+
+    def jpeg_worker(self):
+        img = self.display_obj._repr_jpeg_()
+        if imghdr.what('', img) == "jpeg":
+            self.results["jpeg"] = True
+
+    def png_worker(self):
+        img = self.display_obj._repr_png_()
+        if imghdr.what('', img) == "png":
+            self.results["png"] = True
+
+    def svg_worker(self):
+        img = self.display_obj._repr_svg_()
+        if img[:4] == '<svg':
+            self.results["svg"] = True
+
+    def pdf_worker(self):
+        img = self.display_obj._repr_pdf_()
+        if img[:4] == '%PDF':
+            self.results["pdf"] = True

plotly/tools.py

@@ -13,6 +13,7 @@
 import os.path
 import warnings
 import six
+import requests
 
 from plotly import utils
 from plotly import exceptions
@@ -30,6 +31,12 @@ def warning_on_one_line(message, category, filename, lineno, file=None, line=Non
 except ImportError:
     _matplotlylib_imported = False
 
+try:
+    import IPython
+    _ipython_imported = True
+except ImportError:
+    _ipython_imported = False
+
 PLOTLY_DIR = os.path.join(os.path.expanduser("~"), ".plotly")
 CREDENTIALS_FILE = os.path.join(PLOTLY_DIR, ".credentials")
 CONFIG_FILE = os.path.join(PLOTLY_DIR, ".config")
@@ -283,11 +290,21 @@ def embed(file_owner_or_url, file_id=None, width="100%", height=525):
         return html(s, hide=False)
     except:
         pass
-    try:
-        from IPython.display import HTML, display
-        display(HTML(s))
-    except:
-        pass
+    if _ipython_imported:
+        if file_id:
+            url = "{plotly_domain}/~{un}/{fid}".format(
+                plotly_domain=get_config_file()['plotly_domain'],
+                un=file_owner_or_url,
+                fid=file_id)
+        else:
+            url = file_owner_or_url
+        return PlotlyDisplay(url)
+    else:
+        warnings.warn(
+            "Looks like you're not using IPython or Sage to embed this plot. "
+            "If you just want the *embed code*, try using `get_embed()` "
+            "instead."
+            "\nQuestions? [email protected]")
 
 
 ### mpl-related tools ###
@@ -532,3 +549,40 @@ def _replace_newline(obj):
         return s
     else:
         return obj  # we return the actual reference... but DON'T mutate.
+
+
+if _ipython_imported:
+    class PlotlyDisplay(IPython.core.display.HTML):
+        """An IPython display object for use with plotly urls
+
+        PlotlyDisplay objects should be instantiated with a url for a plot.
+        IPython will *choose* the proper display representation from any
+        Python object, and using provided methods if they exist. By defining
+        the following, if an HTML display is unusable, the PlotlyDisplay
+        object can provide alternate representations.
+
+        """
+        def __init__(self, url):
+            self.resource = url
+            self.embed_code = get_embed(url)
+            super(PlotlyDisplay, self).__init__(data=self.embed_code)
+
+        def _repr_svg_(self):
+            url = self.resource + ".svg"
+            res = requests.get(url)
+            return res.content
+
+        def _repr_png_(self):
+            url = self.resource + ".png"
+            res = requests.get(url)
+            return res.content
+
+        def _repr_pdf_(self):
+            url = self.resource + ".pdf"
+            res = requests.get(url)
+            return res.content
+
+        def _repr_jpeg_(self):
+            url = self.resource + ".jpeg"
+            res = requests.get(url)
+            return res.content