Merge pull request #21 from eriknw/fix_ipywidgets

Fix displaying of outputs

afar/_magic.py → afar/_abra.py

@@ -81,7 +81,7 @@ def __setstate__(self, state):
         self._scoped = scoped_function(func, outer_scope)
 
 
-def abracadabra(runner):
+def cadabra(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)

afar/_core.py

@@ -5,8 +5,8 @@
 
 from dask import distributed
 
-from ._magic import abracadabra
-from ._printing import RecordPrint, print_outputs, print_outputs_async
+from ._abra import cadabra
+from ._printing import PrintRecorder, print_outputs, print_outputs_async
 from ._reprs import repr_afar
 from ._utils import is_kernel, supports_async_output
 from ._where import find_where
@@ -185,7 +185,7 @@ def _exit(self, exc_type, exc_value, exc_traceback):
             endline = maxline + 5  # give us some wiggle room
 
         self.context_body = get_body(self._lines[self._body_start : endline])
-        self._magic_func, names, futures = abracadabra(self)
+        self._magic_func, names, futures = cadabra(self)
         display_expr = self._magic_func._display_expr
 
         if self._where == "remotely":
@@ -267,9 +267,7 @@ def _exit(self, exc_type, exc_value, exc_traceback):
 
                 out = Output()
                 display(out)
-                with out:
-                    print("\N{SPARKLES} Running afar... \N{SPARKLES}")
-                # Can we show `distributed.progress` right here?
+                out.append_stdout("\N{SPARKLES} Running afar... \N{SPARKLES}")
                 stdout_future.add_done_callback(
                     partial(print_outputs_async, out, stderr_future, repr_future)
                 )
@@ -316,7 +314,7 @@ class Get(Run):
 
 def run_afar(magic_func, names, futures, capture_print):
     if capture_print:
-        rec = RecordPrint()
+        rec = PrintRecorder()
         if "print" in magic_func._scoped.builtin_names and "print" not in futures:
             sfunc = magic_func._scoped.bind(futures, print=rec)
         else:

afar/_printing.py

@@ -2,6 +2,7 @@
 import sys
 from io import StringIO
 from threading import Lock, local
+from time import sleep
 
 from ._reprs import display_repr
 
@@ -16,7 +17,7 @@ def __call__(self, *args, **kwargs):
         return self.printer(*args, **kwargs)
 
 
-class RecordPrint:
+class PrintRecorder:
     n = 0
     local_print = LocalPrint()
     print_lock = Lock()
@@ -27,17 +28,17 @@ def __init__(self):
 
     def __enter__(self):
         with self.print_lock:
-            if RecordPrint.n == 0:
+            if PrintRecorder.n == 0:
                 LocalPrint.printer = builtins.print
                 builtins.print = self.local_print
-            RecordPrint.n += 1
+            PrintRecorder.n += 1
         self.local_print.printer = self
         return self
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
         with self.print_lock:
-            RecordPrint.n -= 1
-            if RecordPrint.n == 0:
+            PrintRecorder.n -= 1
+            if PrintRecorder.n == 0:
                 builtins.print = LocalPrint.printer
         self.local_print.printer = LocalPrint.printer
         return False
@@ -72,18 +73,30 @@ def print_outputs_async(out, stderr_future, repr_future, stdout_future):
 
     This is used as a callback to `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
-    out.clear_output()
-    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)
+    try:
+        stdout_val = stdout_future.result()
+        out.clear_output()
+        count = 0
+        while out.outputs:
+            # See: https://github.com/jupyter-widgets/ipywidgets/issues/3260
+            count += 1
+            if count == 100:  # 0.5 seconds
+                # This doesn't appear to always clear correctly in JupyterLab.
+                # I don't know why.  I'm still investigating.
+                out.outputs = type(out.outputs)()  # is this safe?
+                break
+            sleep(0.005)
+        if stdout_val:
+            out.append_stdout(stdout_val)
+
+        stderr_val = stderr_future.result()
+        if stderr_val:
+            out.append_stderr(stderr_val)
+
+        if repr_future is not None:
+            repr_val = repr_future.result()
             if repr_val is not None:
-                display_repr(repr_val)
+                display_repr(repr_val, out=out)
+    except Exception as exc:
+        print(exc, file=sys.stderr)
+        raise

afar/_reprs.py

@@ -84,18 +84,28 @@ def __repr__(self):
         return self.val
 
 
-def display_repr(results):
+def display_repr(results, out=None):
     """Display results from `repr_afar` locally in IPython/Jupyter"""
     val, method_name, is_exception = results
     if is_exception:
-        print(val, file=sys.stderr)
+        if out is None:
+            print(val, file=sys.stderr)
+        else:
+            out.append_stderr(val)
         return
     if val is None and method_name is None:
         return
+
+    from IPython.display import display
+
     if method_name == "_ipython_display_":
-        val._ipython_display_()
+        if out is None:
+            display(val)
+        else:
+            out.append_display_data(val)
     else:
-        from IPython.display import display
-
         mimic = MimicRepr(val, method_name)
-        display(mimic)
+        if out is None:
+            display(mimic)
+        else:
+            out.append_display_data(mimic)