def capture_output(
stdout: bool = True,
stderr: bool = True,
display: bool = True
) -> typing.Iterable[typing.Sequence[typing.Dict[str, typing.Any]]]:
_stdout = sys.stdout
_stderr = sys.stderr
_display_hook = sys.displayhook
shell = get_ipython()
_display_pub = shell.display_pub if shell else None
outputs = _CaptureOutputs()
try:
if stdout:
sys.stdout = _CaptureWriter('stdout', outputs)
if stderr:
sys.stderr = _CaptureWriter('stderr', outputs)
if display and shell:
shell.display_pub = CapturingDisplayPublisher()
shell.display_pub.outputs = outputs
sys.displayhook = CapturingDisplayHook(shell, outputs)
yield outputs
finally:
sys.stdout = _stdout
sys.stderr = _stderr
sys.displayhook = _display_hook
if shell:
shell.display_pub = _display_pub
def __enter__(self):
from IPython.core.getipython import get_ipython
from IPython.core.displaypub import CapturingDisplayPublisher
self.sys_stdout = sys.stdout
self.sys_stderr = sys.stderr
if self.display:
self.shell = get_ipython()
if self.shell is None:
self.save_display_pub = None
self.display = False
stdout = stderr = outputs = None
if self.stdout:
stdout = sys.stdout = StringIO()
if self.stderr:
stderr = sys.stderr = StringIO()
if self.display:
self.save_display_pub = self.shell.display_pub
self.shell.display_pub = CapturingDisplayPublisher()
outputs = self.shell.display_pub.outputs
return CapturedIO(stdout, stderr, outputs)
def _handle_biyam(self, stream: ZMQStream, ident: typing.List[bytes],
msg: typing.Dict[str, typing.Any]) -> None:
content = msg.get('content')
caller = content.get('caller', None)
if caller is None:
return
comm = _CommWriter(self._kernel, caller)
stdout_ = sys.stdout
stderr_ = sys.stderr
displayhook_ = sys.displayhook
shell = self._shell
display_pub_ = shell.display_pub if shell else None
try:
sys.stdout = typing.cast(typing.Any, _PseudoWriter(comm, 'stdout'))
sys.stderr = typing.cast(typing.Any, _PseudoWriter(comm, 'stderr'))
if shell:
outputs = _DummyList(comm)
shell.display_pub = CapturingDisplayPublisher()
shell.display_pub.outputs = outputs
sys.displayhook = CapturingDisplayHook(shell, outputs)
mod_name = content.get('mod', '__main__')
mod = importlib.import_module(mod_name)
fn_ = getattr(mod, content.get('fn', 'main'))
if not hasattr(fn_, '_biyam_meta_'):
raise PermissionError(fn_.__name__ + ' is not exposed')
ret = fn_(**content.get('args', {}))
comm.resolve(ret)
except Exception:
traceback.print_exc()
comm.reject()
finally:
sys.stdout = stdout_
sys.stderr = stderr_
sys.displayhook = displayhook_
if shell:
shell.display_pub = display_pub_
