Merge pull request #1946 from tkf/terminal-image-handler

Add image message handler in ZMQTerminalInteractiveShell

This change introduces several handlers for messages which contain image in
ZMQTerminalInteractiveShell. This is useful, for example, when
connecting to the kernel in which pylab inline backend is activated.

This PR will fix  #1575.

IPython/frontend/terminal/console/tests/test_image_handler.py

@@ -0,0 +1,95
+#-----------------------------------------------------------------------------
+# Copyright (C) 2012 The IPython Development Team
+#
+# Distributed under the terms of the BSD License. The full license is in
+# the file COPYING, distributed as part of this software.
+#-----------------------------------------------------------------------------
+
+import os
+import sys
+import unittest
+import base64
+
+from IPython.zmq.kernelmanager import KernelManager
+from IPython.frontend.terminal.console.interactiveshell \
+    import ZMQTerminalInteractiveShell
+from IPython.utils.tempdir import TemporaryDirectory
+from IPython.testing.tools import monkeypatch
+from IPython.testing.decorators import skip_without
+from IPython.utils.ipstruct import Struct
+
+
+SCRIPT_PATH = os.path.join(
+    os.path.abspath(os.path.dirname(__file__)), 'writetofile.py')
+
+
+class ZMQTerminalInteractiveShellTestCase(unittest.TestCase):
+
+    def setUp(self):
+        km = KernelManager()
+        self.shell = ZMQTerminalInteractiveShell(kernel_manager=km)
+        self.raw = b'dummy data'
+        self.mime = 'image/png'
+        self.data = {self.mime: base64.encodestring(self.raw).decode('ascii')}
+
+    def test_no_call_by_default(self):
+        def raise_if_called(*args, **kwds):
+            assert False
+
+        shell = self.shell
+        shell.handle_image_PIL = raise_if_called
+        shell.handle_image_stream = raise_if_called
+        shell.handle_image_tempfile = raise_if_called
+        shell.handle_image_callable = raise_if_called
+
+        shell.handle_image(None, None)  # arguments are dummy
+
+    @skip_without('PIL')
+    def test_handle_image_PIL(self):
+        import PIL.Image
+
+        open_called_with = []
+        show_called_with = []
+
+        def fake_open(arg):
+            open_called_with.append(arg)
+            return Struct(show=lambda: show_called_with.append(None))
+
+        with monkeypatch(PIL.Image, 'open', fake_open):
+            self.shell.handle_image_PIL(self.data, self.mime)
+
+        self.assertEqual(len(open_called_with), 1)
+        self.assertEqual(len(show_called_with), 1)
+        self.assertEqual(open_called_with[0].getvalue(), self.raw)
+
+    def check_handler_with_file(self, inpath, handler):
+        shell = self.shell
+        configname = '{0}_image_handler'.format(handler)
+        funcname = 'handle_image_{0}'.format(handler)
+
+        assert hasattr(shell, configname)
+        assert hasattr(shell, funcname)
+
+        with TemporaryDirectory() as tmpdir:
+            outpath = os.path.join(tmpdir, 'data')
+            cmd = [sys.executable, SCRIPT_PATH, inpath, outpath]
+            setattr(shell, configname, cmd)
+            getattr(shell, funcname)(self.data, self.mime)
+            # cmd is called and file is closed.  So it's safe to open now.
+            with open(outpath, 'rb') as file:
+                transferred = file.read()
+
+        self.assertEqual(transferred, self.raw)
+
+    def test_handle_image_stream(self):
+        self.check_handler_with_file('-', 'stream')
+
+    def test_handle_image_tempfile(self):
+        self.check_handler_with_file('{file}', 'tempfile')
+
+    def test_handle_image_callable(self):
+        called_with = []
+        self.shell.callable_image_handler = called_with.append
+        self.shell.handle_image_callable(self.data, self.mime)
+        self.assertEqual(len(called_with), 1)
+        assert called_with[0] is self.data

IPython/frontend/terminal/console/tests/writetofile.py

@@ -0,0 +1,33
+#-----------------------------------------------------------------------------
+# Copyright (C) 2012 The IPython Development Team
+#
+# Distributed under the terms of the BSD License. The full license is in
+# the file COPYING, distributed as part of this software.
+#-----------------------------------------------------------------------------
+
+"""
+Copy data from input file to output file for testing.
+
+Command line usage:
+
+    python writetofile.py INPUT OUTPUT
+
+Binary data from INPUT file is copied to OUTPUT file.
+If INPUT is '-', stdin is used.
+
+"""
+
+if __name__ == '__main__':
+    import sys
+    from IPython.utils.py3compat import PY3
+    (inpath, outpath) = sys.argv[1:]
+
+    if inpath == '-':
+        if PY3:
+            infile = sys.stdin.buffer
+        else:
+            infile = sys.stdin
+    else:
+        infile = open(inpath, 'rb')
+
+    open(outpath, 'w+b').write(infile.read())

IPython/utils/tests/test_tempdir.py

@@ -0,0 +1,20
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2012-  The IPython Development Team
+#
+#  Distributed under the terms of the BSD License.  The full license is in
+#  the file COPYING, distributed as part of this software.
+#-----------------------------------------------------------------------------
+
+import os
+
+from IPython.utils.tempdir import NamedFileInTemporaryDirectory
+
+
+def test_named_file_in_temporary_directory():
+    with NamedFileInTemporaryDirectory('filename') as file:
+        name = file.name
+        assert not file.closed
+        assert os.path.exists(name)
+        file.write(b'test')
+    assert file.closed
+    assert not os.path.exists(name)

IPython/frontend/terminal/console/interactiveshell.py

@@ -19,15 +19,26 from __future__ import print_function
 
 import bdb
 import signal
+import os
 import sys
 import time
+import subprocess
+from io import BytesIO
+import base64
 
 from Queue import Empty
 
+try:
+    from contextlib import nested
+except:
+    from IPython.utils.nested_context import nested
+
 from IPython.core.alias import AliasManager, AliasError
 from IPython.core import page
 from IPython.utils.warn import warn, error, fatal
 from IPython.utils import io
+from IPython.utils.traitlets import List, Enum, Any
+from IPython.utils.tempdir import NamedFileInTemporaryDirectory
 
 from IPython.frontend.terminal.interactiveshell import TerminalInteractiveShell
 from IPython.frontend.terminal.console.completer import ZMQCompleter
@@ -36,7 +47,64 from IPython.frontend.terminal.console.completer import ZMQCompleter
 class ZMQTerminalInteractiveShell(TerminalInteractiveShell):
     """A subclass of TerminalInteractiveShell that uses the 0MQ kernel"""
     _executing = False
-    
+
+    image_handler = Enum(('PIL', 'stream', 'tempfile', 'callable'),
+                         config=True, help=
+        """
+        Handler for image type output.  This is useful, for example,
+        when connecting to the kernel in which pylab inline backend is
+        activated.  There are four handlers defined.  'PIL': Use
+        Python Imaging Library to popup image; 'stream': Use an
+        external program to show the image.  Image will be fed into
+        the STDIN of the program.  You will need to configure
+        `stream_image_handler`; 'tempfile': Use an external program to
+        show the image.  Image will be saved in a temporally file and
+        the program is called with the temporally file.  You will need
+        to configure `tempfile_image_handler`; 'callable': You can set
+        any Python callable which is called with the image data.  You
+        will need to configure `callable_image_handler`.
+        """
+    )
+
+    stream_image_handler = List(config=True, help=
+        """
+        Command to invoke an image viewer program when you are using
+        'stream' image handler.  This option is a list of string where
+        the first element is the command itself and reminders are the
+        options for the command.  Raw image data is given as STDIN to
+        the program.
+        """
+    )
+
+    tempfile_image_handler = List(config=True, help=
+        """
+        Command to invoke an image viewer program when you are using
+        'tempfile' image handler.  This option is a list of string
+        where the first element is the command itself and reminders
+        are the options for the command.  You can use {file} and
+        {format} in the string to represent the location of the
+        generated image file and image format.
+        """
+    )
+
+    callable_image_handler = Any(config=True, help=
+        """
+        Callable object called via 'callable' image handler with one
+        argument, `data`, which is `msg["content"]["data"]` where
+        `msg` is the message from iopub channel.  For exmaple, you can
+        find base64 encoded PNG data as `data['image/png']`.
+        """
+    )
+
+    mime_preference = List(
+        default_value=['image/png', 'image/jpeg', 'image/svg+xml'],
+        config=True, allow_none=False, help=
+        """
+        Preferred object representation MIME type in order.  First
+        matched MIME type will be used.
+        """
+    )
+
     def __init__(self, *args, **kwargs):
         self.km = kwargs.pop('kernel_manager')
         self.session_id = self.km.session.session
@@ -163,6 +231,7 class ZMQTerminalInteractiveShell(TerminalInteractiveShell):
                 elif msg_type == 'pyout':
                     self.execution_count = int(sub_msg["content"]["execution_count"])
                     format_dict = sub_msg["content"]["data"]
+                    self.handle_rich_data(format_dict)
                     # taken from DisplayHook.__call__:
                     hook = self.displayhook
                     hook.start_displayhook()
@@ -171,6 +240,61 class ZMQTerminalInteractiveShell(TerminalInteractiveShell):
                     hook.log_output(format_dict)
                     hook.finish_displayhook()
 
+                elif msg_type == 'display_data':
+                    self.handle_rich_data(sub_msg["content"]["data"])
+
+    _imagemime = {
+        'image/png': 'png',
+        'image/jpeg': 'jpeg',
+        'image/svg+xml': 'svg',
+    }
+
+    def handle_rich_data(self, data):
+        for mime in self.mime_preference:
+            if mime in data and mime in self._imagemime:
+                self.handle_image(data, mime)
+                return
+
+    def handle_image(self, data, mime):
+        handler = getattr(
+            self, 'handle_image_{0}'.format(self.image_handler), None)
+        if handler:
+            handler(data, mime)
+
+    def handle_image_PIL(self, data, mime):
+        if mime not in ('image/png', 'image/jpeg'):
+            return
+        import PIL.Image
+        raw = base64.decodestring(data[mime].encode('ascii'))
+        img = PIL.Image.open(BytesIO(raw))
+        img.show()
+
+    def handle_image_stream(self, data, mime):
+        raw = base64.decodestring(data[mime].encode('ascii'))
+        imageformat = self._imagemime[mime]
+        fmt = dict(format=imageformat)
+        args = [s.format(**fmt) for s in self.stream_image_handler]
+        with open(os.devnull, 'w') as devnull:
+            proc = subprocess.Popen(
+                args, stdin=subprocess.PIPE,
+                stdout=devnull, stderr=devnull)
+            proc.communicate(raw)
+
+    def handle_image_tempfile(self, data, mime):
+        raw = base64.decodestring(data[mime].encode('ascii'))
+        imageformat = self._imagemime[mime]
+        filename = 'tmp.{0}'.format(imageformat)
+        with nested(NamedFileInTemporaryDirectory(filename),
+                    open(os.devnull, 'w')) as (f, devnull):
+            f.write(raw)
+            f.flush()
+            fmt = dict(file=f.name, format=imageformat)
+            args = [s.format(**fmt) for s in self.tempfile_image_handler]
+            subprocess.call(args, stdout=devnull, stderr=devnull)
+
+    def handle_image_callable(self, data, mime):
+        self.callable_image_handler(data)
+
     def handle_stdin_request(self, timeout=0.1):
         """ Method to capture raw_input
         """

IPython/utils/tempdir.py

@@ -3,14 +3,14
 This is copied from the stdlib and will be standard in Python 3.2 and onwards.
 """
 
+import os as _os
+
 # This code should only be used in Python versions < 3.2, since after that we
 # can rely on the stdlib itself.
 try:
     from tempfile import TemporaryDirectory
 
 except ImportError:
-    
-    import os as _os
     from tempfile import mkdtemp, template
 
     class TemporaryDirectory(object):
@@ -74,3 +74,33 except ImportError:
                 self._rmdir(path)
             except self._os_error:
                 pass
+
+
+class NamedFileInTemporaryDirectory(object):
+
+    def __init__(self, filename, mode='w+b', bufsize=-1, **kwds):
+        """
+        Open a file named `filename` in a temporary directory.
+
+        This context manager is preferred over `NamedTemporaryFile` in
+        stdlib `tempfile` when one needs to reopen the file.
+
+        Arguments `mode` and `bufsize` are passed to `open`.
+        Rest of the arguments are passed to `TemporaryDirectory`.
+
+        """
+        self._tmpdir = TemporaryDirectory(**kwds)
+        path = _os.path.join(self._tmpdir.name, filename)
+        self.file = open(path, mode, bufsize)
+
+    def cleanup(self):
+        self.file.close()
+        self._tmpdir.cleanup()
+
+    __del__ = cleanup
+
+    def __enter__(self):
+        return self.file
+
+    def __exit__(self, type, value, traceback):
+        self.cleanup()