diff --git a/afar/_magic.py b/afar/_abra.py similarity index 99% rename from afar/_magic.py rename to afar/_abra.py index 24c3d8e..8aed5bb 100644 --- a/afar/_magic.py +++ b/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) diff --git a/afar/_core.py b/afar/_core.py index 6da8858..2f2cdf5 100644 --- a/afar/_core.py +++ b/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: diff --git a/afar/_printing.py b/afar/_printing.py index 3059c1b..465870f 100644 --- a/afar/_printing.py +++ b/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 diff --git a/afar/_reprs.py b/afar/_reprs.py index cb97064..19ad0f9 100644 --- a/afar/_reprs.py +++ b/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)