From ee83d92c83082d0ba05be2f50612c703e108c116 2012-09-29 12:53:43 From: Bussonnier Matthias <bussonniermatthias@gmail.com> Date: 2012-09-29 12:53:43 Subject: [PATCH] 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. --- diff --git a/IPython/frontend/terminal/console/interactiveshell.py b/IPython/frontend/terminal/console/interactiveshell.py index 1291822..983d154 100644 --- a/IPython/frontend/terminal/console/interactiveshell.py +++ b/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 """ diff --git a/IPython/frontend/terminal/console/tests/test_image_handler.py b/IPython/frontend/terminal/console/tests/test_image_handler.py new file mode 100644 index 0000000..b10df00 --- /dev/null +++ b/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 diff --git a/IPython/frontend/terminal/console/tests/writetofile.py b/IPython/frontend/terminal/console/tests/writetofile.py new file mode 100644 index 0000000..a6f9e28 --- /dev/null +++ b/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()) diff --git a/IPython/utils/tempdir.py b/IPython/utils/tempdir.py index 57f02fc..364e119 100644 --- a/IPython/utils/tempdir.py +++ b/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() diff --git a/IPython/utils/tests/test_tempdir.py b/IPython/utils/tests/test_tempdir.py new file mode 100644 index 0000000..8e919e5 --- /dev/null +++ b/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)