Add image message handler in ZMQTerminalInteractiveShell

IPython/frontend/terminal/console/interactiveshell.py

@@ -19,15 +19,26 @@
 
 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 @@
 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 @@ def handle_iopub(self):
                 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 @@ def handle_iopub(self):
                     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/frontend/terminal/console/tests/test_image_handler.py

@@ -0,0 +1,92 @@
+#-----------------------------------------------------------------------------
+# 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 unittest
+import base64
+
+from IPython.zmq.kernelmanager import KernelManager
+from IPython.frontend.terminal.console.interactiveshell \
+    import ZMQTerminalInteractiveShell
+from IPython.utils.tempdir import NamedFileInTemporaryDirectory
+from IPython.testing.tools import monkeypatch
+from IPython.testing.decorators import skip_without
+from IPython.utils.ipstruct import Struct
+from IPython.utils.process import find_cmd
+
+
+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
+        shell.handle_image_stream
+        shell.handle_image_tempfile
+        shell.handle_image_callable
+
+        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)
+
+        assert len(open_called_with) == 1
+        assert len(show_called_with) == 1
+        assert 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 NamedFileInTemporaryDirectory('data') as file:
+            cmd = [find_cmd('python'), SCRIPT_PATH, inpath, file.name]
+            setattr(shell, configname, cmd)
+            getattr(shell, funcname)(self.data, self.mime)
+            transferred = file.read()
+
+        assert 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)
+        assert 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/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 @@ def _rmtree(self, path):
                 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()

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)