Capture print statements and show them in notebooks with ipywidgets.

This seems to work well in Jupyteer Notebook, but not Jupyter Lab.
Also, unlike before, stderr isn't streamed (how does that work btw?).
So, all output is sent back to the client upon the work being completed.

Author: eriknw

afar/core.py

@@ -3,13 +3,26 @@
 import inspect
 import io
 import sys
+import threading
 import weakref
+from functools import partial
 
 import innerscope
 from dask import distributed
 
 from . import reprs
 
+
+def _supports_async_output():
+    if reprs.is_kernel() and not reprs.in_terminal():
+        try:
+            import ipywidgets  # noqa
+        except ImportError:
+            return False
+        return True
+    return False
+
+
 _errors_to_locations = {}
 try:
     remotely
@@ -137,11 +150,13 @@ def __enter__(self):
             # Try to fine the source if we are in %%time or %%timeit magic
             if (
                 self._frame.f_code.co_filename in {"<timed exec>", "<magic-timeit>"}
-                and reprs.in_ipython()
+                and reprs.is_kernel()
             ):
-                import IPython
+                from IPython import get_ipython
 
-                ip = IPython.get_ipython()
+                ip = get_ipython()
+                if ip is None:
+                    raise
                 cell = ip.history_manager._i00  # The current cell!
                 lines = cell.splitlines(keepends=True)
                 # strip the magic
@@ -233,6 +248,7 @@ def _exit(self, exc_type, exc_value, exc_traceback):
         self.context_body = get_body(self._lines[self._body_start : endline])
         self._magic_func, names, futures = abracadabra(self)
         display_expr = self._magic_func._display_expr
+        has_print = "print" in self._magic_func._scoped.builtin_names
 
         if self._where == "remotely":
             if client is None:
@@ -268,16 +284,19 @@ def _exit(self, exc_type, exc_value, exc_traceback):
             weak_futures.add(remote_dict)
             magic_func.release()  # Let go ASAP
             if display_expr:
-                repr_val = client.submit(
+                repr_future = client.submit(
                     reprs.repr_afar,
                     client.submit(get_afar, remote_dict, "_afar_return_value_"),
                     self._magic_func._repr_methods,
                 )
-                weak_futures.add(repr_val)
-            stdout_future = client.submit(get_afar, remote_dict, "_afar_stdout_")
-            weak_futures.add(stdout_future)
-            stderr_future = client.submit(get_afar, remote_dict, "_afar_stderr_")
-            weak_futures.add(stderr_future)
+                weak_futures.add(repr_future)
+            else:
+                repr_future = None
+            if display_expr or has_print or _supports_async_output():
+                stdout_future = client.submit(get_afar, remote_dict, "_afar_stdout_")
+                weak_futures.add(stdout_future)
+                stderr_future = client.submit(get_afar, remote_dict, "_afar_stderr_")
+                weak_futures.add(stderr_future)
             if self._gather_data:
                 futures_to_name = {
                     client.submit(get_afar, remote_dict, name, **submit_kwargs): name
@@ -294,25 +313,41 @@ def _exit(self, exc_type, exc_value, exc_traceback):
                     self.data[name] = future
                 remote_dict.release()  # Let go ASAP
 
-            # blocks!
-            stdout_val = stdout_future.result()
-            if stdout_val:
-                print(stdout_val, end="")
-            stdout_future.release()
-            stderr_val = stderr_future.result()
-            if stderr_val:
-                print(stderr_val, end="", file=sys.stderr)
-            stderr_future.release()
-            if display_expr:
-                reprs.display_repr(repr_val.result())  # This blocks!
-                repr_val.release()
+            if _supports_async_output():
+                # Display in `out` cell when data is ready: non-blocking
+                from IPython.display import display
+                from ipywidgets import Output
+
+                out = Output()
+                display(out)
+                # Can we show `distributed.progress` right here?
+                stdout_future.add_done_callback(
+                    partial(_display_outputs, out, stderr_future, repr_future)
+                )
+            elif display_expr or has_print:
+                # blocks!
+                stdout_val = stdout_future.result()
+                stdout_future.release()
+                if stdout_val:
+                    print(stdout_val, end="")
+                stderr_val = stderr_future.result()
+                stderr_future.release()
+                if stderr_val:
+                    print(stderr_val, end="", file=sys.stderr)
+                if display_expr:
+                    repr_val = repr_future.result()
+                    repr_future.release()
+                    if repr_val is not None:
+                        reprs.display_repr(repr_val)
         elif self._where == "locally":
             # Run locally.  This is handy for testing and debugging.
             results = self._magic_func()
             for name in names:
                 self.data[name] = results[name]
             if display_expr:
-                reprs.IPython.display.display(results.return_value)
+                from IPython.dislpay import display
+
+                display(results.return_value)
         elif self._where == "later":
             return True
         else:
@@ -342,11 +377,28 @@ class Get(Run):
     _gather_data = True
 
 
+def _display_outputs(out, stderr_future, repr_future, stdout_future):
+    stdout_val = stdout_future.result()
+    stderr_val = stderr_future.result()
+    if repr_future is not None:
+        repr_val = repr_future.result()
+    else:
+        repr_val = None
+    if stdout_val or stderr_val or repr_val is not None:
+        with out:
+            if stdout_val:
+                print(stdout_val, end="")
+            if stderr_val:
+                print(stderr_val, end="", file=sys.stderr)
+            if repr_val is not None:
+                reprs.display_repr(repr_val)
+
+
 def abracadabra(runner):
     # Create a new function from the code block of the context.
     # For now, we require that the source code is available.
     source = "def _afar_magic_():\n" + "".join(runner.context_body)
-    func, display_expr = create_func(source, runner._frame.f_globals, reprs.in_ipython())
+    func, display_expr = create_func(source, runner._frame.f_globals, reprs.is_kernel())
 
     # If no variable names were given, only get the last assignment
     names = runner.names
@@ -422,26 +474,48 @@ def __setstate__(self, state):
         self._scoped = innerscope.scoped_function(func, outer_scope)
 
 
+# Here's the plan: we'll capture all print statements to stdout and stderr
+# on the current thread.  But, we need to leave the other threads alone!
+# So, use `threading.local` and a lock for some ugly capturing.
+class LocalPrint(threading.local):
+    printer = None
+
+    def __call__(self, *args, **kwargs):
+        return self.printer(*args, **kwargs)
+
+
 class RecordPrint:
+    n = 0
+    local_print = LocalPrint()
+    print_lock = threading.Lock()
+
     def __init__(self):
         self.stdout = io.StringIO()
         self.stderr = io.StringIO()
 
     def __enter__(self):
-        self.print = builtins.print
-        builtins.print = self
+        with self.print_lock:
+            if RecordPrint.n == 0:
+                LocalPrint.printer = builtins.print
+                builtins.print = self.local_print
+            RecordPrint.n += 1
+        self.local_print.printer = self
         return self
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
-        builtins.print = self.print
+        with self.print_lock:
+            RecordPrint.n -= 1
+            if RecordPrint.n == 0:
+                builtins.print = LocalPrint.printer
+        self.local_print.printer = LocalPrint.printer
         return False
 
     def __call__(self, *args, file=None, **kwargs):
         if file is None or file is sys.stdout:
             file = self.stdout
         elif file is sys.stderr:
             file = self.stderr
-        self.print(*args, **kwargs, file=file)
+        LocalPrint.printer(*args, **kwargs, file=file)
 
 
 def run_afar(magic_func, names, futures):

afar/reprs.py

@@ -3,17 +3,15 @@
 import traceback
 from types import CodeType, FunctionType
 
-try:
-    import IPython
+from distributed.utils import is_kernel  # noqa
 
-    def in_ipython():
-        return IPython.get_ipython() is not None
 
-
-except ImportError:
-
-    def in_ipython():
+def in_terminal():
+    if "IPython" not in sys.modules:  # IPython hasn't been imported
         return False
+    from IPython import get_ipython
+
+    return type(get_ipython()).__name__ == "TerminalInteractiveShell"
 
 
 if hasattr(CodeType, "replace"):
@@ -88,7 +86,9 @@ def __getattr__(self, attr):
 
 def get_repr_methods():
     """List of repr methods that IPython/Jupyter tries to use"""
-    ip = IPython.get_ipython()
+    from IPython import get_ipython
+
+    ip = get_ipython()
     if ip is None:
         return
     attr_recorder = AttrRecorder()
@@ -101,6 +101,8 @@ def repr_afar(val, repr_methods):
 
     We call this on a remote object.
     """
+    if val is None:
+        return None
     for method_name in repr_methods:
         method = getattr(val, method_name, None)
         if method is None:
@@ -134,11 +136,15 @@ def display_repr(results):
     if is_exception:
         print(val, file=sys.stderr)
         return
+    if val is None and method_name is None:
+        return
     if method_name == "_ipython_display_":
         val._ipython_display_()
     else:
+        from IPython.display import display
+
         mimic = MimicRepr(val, method_name)
-        IPython.display.display(mimic)
+        display(mimic)
 
 
 class MimicRepr: