PR: [alternative to #438] Update matplotlib figures while waiting for input

ipykernel/kernelbase.py

@@ -13,6 +13,7 @@
 import sys
 import time
 import uuid
+import threading
 
 try:
     # jupyter_client >= 5, use tz-aware now
@@ -30,7 +31,7 @@
 from traitlets.config.configurable import SingletonConfigurable
 from IPython.core.error import StdinNotImplementedError
 from ipython_genutils import py3compat
-from ipython_genutils.py3compat import unicode_type, string_types
+from ipython_genutils.py3compat import unicode_type, string_types, PY3
 from ipykernel.jsonutil import json_clean
 from traitlets import (
     Any, Instance, Float, Dict, List, Set, Integer, Unicode, Bool,
@@ -41,6 +42,7 @@
 
 from ._version import kernel_protocol_version
 
+
 CONTROL_PRIORITY = 1
 SHELL_PRIORITY = 10
 ABORT_PRIORITY = 20
@@ -170,6 +172,9 @@ def __init__(self, **kwargs):
         for msg_type in self.control_msg_types:
             self.control_handlers[msg_type] = getattr(self, msg_type)
 
+        self._stdin_lock = threading.Lock()
+        self._enable_input_matplotlib_processing = True
+
     @gen.coroutine
     def dispatch_control(self, msg):
         """dispatch control requests"""
@@ -860,6 +865,7 @@ def raw_input(self, prompt=''):
         )
 
     def _input_request(self, prompt, ident, parent, password=False):
+        """Send an input request to the frontend and wait for the reply."""
         # Flush output before making the request.
         sys.stderr.flush()
         sys.stdout.flush()
@@ -873,22 +879,14 @@ def _input_request(self, prompt, ident, parent, password=False):
                 else:
                     raise
 
-        # Send the input request.
-        content = json_clean(dict(prompt=prompt, password=password))
-        self.session.send(self.stdin_socket, u'input_request', content, parent,
-                          ident=ident)
+        with self._stdin_lock:
+            # Send the input request.
+            content = json_clean(dict(prompt=prompt, password=password))
+            self.session.send(self.stdin_socket, u'input_request',
+                              content, parent, ident=ident)
+            # Await a response.
+            reply = self._wait_input_request_reply()
 
-        # Await a response.
-        while True:
-            try:
-                ident, reply = self.session.recv(self.stdin_socket, 0)
-            except Exception:
-                self.log.warning("Invalid Message:", exc_info=True)
-            except KeyboardInterrupt:
-                # re-raise KeyboardInterrupt, to truncate traceback
-                raise KeyboardInterrupt
-            else:
-                break
         try:
             value = py3compat.unicode_to_str(reply['content']['value'])
         except:
@@ -899,6 +897,59 @@ def _input_request(self, prompt, ident, parent, password=False):
             raise EOFError
         return value
 
+    def _wait_input_request_reply(self):
+        """Wait for an input request reply.
+
+        Raises
+        ------
+        KeyboardInterrupt if a keyboard interrupt is recieved.
+        """
+        # Await a response.
+        while True:
+            try:
+                # Try a non blocking recv
+                ident, reply = self.session.recv(
+                    self.stdin_socket, zmq.NOBLOCK)
+                if reply:
+                    return reply
+                if not self._input_request_loop_step():
+                    # Wait until a reply is recieved
+                    ident, reply = self.session.recv(self.stdin_socket, 0)
+                    return reply
+            except Exception:
+                self.log.warning("Invalid Message:", exc_info=True)
+            except KeyboardInterrupt:
+                # re-raise KeyboardInterrupt, to truncate traceback
+                raise KeyboardInterrupt
+
+    def _input_request_loop_step(self):
+        """
+        Do one step of the input request loop.
+
+        Returns False if no additional steps are needed, otherwise True.
+        """
+        # Allow matplotlib figures using GUI frameworks (e.g. qt, wx, gtk, tk)
+        # to update
+        if PY3:
+            is_main_thread = (
+                threading.current_thread() is threading.main_thread())
+        else:
+            # Python 2
+            is_main_thread = isinstance(
+                threading.current_thread(), threading._MainThread)
+
+        if (is_main_thread and 'matplotlib.pyplot' in sys.modules and
+                self._enable_input_matplotlib_processing):
+            # matplotlib needs to be imported after app.launch_new_instance()
+            import matplotlib.pyplot as plt
+            if plt.get_fignums():
+                # If there is a figure, update it, wait, and run the next step.
+                plt.gcf().canvas.flush_events()
+                time.sleep(0.01)
+                return True
+        # No more steps are needed
+        return False
+
     def _at_shutdown(self):
         """Actions taken at shutdown by the kernel, called by python's atexit.
         """