From 70c6441688927da456ec6f7fba432ee88a5da63a Mon Sep 17 00:00:00 2001
From: Peter <cHYzZQo@netflix.com>
Date: Mon, 12 Mar 2018 18:39:14 -0700
Subject: [PATCH] kernelapp: Preserve stdout and stderr in kernel

---
 ipykernel/iostream.py  | 26 +++++++++++++++++++++++++-
 ipykernel/kernelapp.py | 13 +++++++++++--
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/ipykernel/iostream.py b/ipykernel/iostream.py
index b1a62557a..7dc519a40 100644
--- a/ipykernel/iostream.py
+++ b/ipykernel/iostream.py
@@ -276,7 +276,7 @@ class OutStream(TextIOBase):
     topic = None
     encoding = 'UTF-8'
 
-    def __init__(self, session, pub_thread, name, pipe=None):
+    def __init__(self, session, pub_thread, name, pipe=None, echo=None):
         if pipe is not None:
             warnings.warn("pipe argument to OutStream is deprecated and ignored",
                 DeprecationWarning)
@@ -296,6 +296,13 @@ def __init__(self, session, pub_thread, name, pipe=None):
         self._flush_pending = False
         self._io_loop = pub_thread.io_loop
         self._new_buffer()
+        self.echo = None
+
+        if echo:
+            if hasattr(echo, 'read') and hasattr(echo, 'write'):
+                self.echo = echo
+            else:
+                raise ValueError("echo argument must be a file like object")
 
     def _is_master_process(self):
         return os.getpid() == self._master_pid
@@ -353,6 +360,15 @@ def _flush(self):
         unless the thread has been destroyed (e.g. forked subprocess).
         """
         self._flush_pending = False
+
+        if self.echo is not None:
+            try:
+                self.echo.flush()
+            except OSError as e:
+                if self.echo is not sys.__stderr__:
+                    print("Flush failed: {}".format(e),
+                          file=sys.__stderr__)
+
         data = self._flush_buffer()
         if data:
             # FIXME: this disables Session's fork-safe check,
@@ -364,6 +380,14 @@ def _flush(self):
                 parent=self.parent_header, ident=self.topic)
 
     def write(self, string):
+        if self.echo is not None:
+            try:
+                self.echo.write(string)
+            except OSError as e:
+                if self.echo is not sys.__stderr__:
+                    print("Write failed: {}".format(e),
+                          file=sys.__stderr__)
+
         if self.pub_thread is None:
             raise ValueError('I/O operation on closed file')
         else:
diff --git a/ipykernel/kernelapp.py b/ipykernel/kernelapp.py
index c729c5369..e0d729e2b 100644
--- a/ipykernel/kernelapp.py
+++ b/ipykernel/kernelapp.py
@@ -139,6 +139,7 @@ def abs_connection_file(self):
     # streams, etc.
     no_stdout = Bool(False, help="redirect stdout to the null device").tag(config=True)
     no_stderr = Bool(False, help="redirect stderr to the null device").tag(config=True)
+    quiet = Bool(True, help="Only send stdout/stderr to output stream").tag(config=True)
     outstream_class = DottedObjectName('ipykernel.iostream.OutStream',
         help="The importstring for the OutStream factory").tag(config=True)
     displayhook_class = DottedObjectName('ipykernel.displayhook.ZMQDisplayHook',
@@ -323,9 +324,17 @@ def init_io(self):
         if self.outstream_class:
             outstream_factory = import_item(str(self.outstream_class))
             sys.stdout.flush()
-            sys.stdout = outstream_factory(self.session, self.iopub_thread, u'stdout')
+
+            e_stdout = None if self.quiet else sys.__stdout__
+            e_stderr = None if self.quiet else sys.__stderr__
+
+            sys.stdout = outstream_factory(self.session, self.iopub_thread,
+                                           u'stdout',
+                                           echo=e_stdout)
             sys.stderr.flush()
-            sys.stderr = outstream_factory(self.session, self.iopub_thread, u'stderr')
+            sys.stderr = outstream_factory(self.session, self.iopub_thread,
+                                           u'stderr',
+                                           echo=e_stderr)
         if self.displayhook_class:
             displayhook_factory = import_item(str(self.displayhook_class))
             self.displayhook = displayhook_factory(self.session, self.iopub_socket)