Avoid echoing onto a captured FD

ipykernel/iostream.py

@@ -364,7 +364,7 @@ def __init__(
         echo : bool
             whether to echo output
         watchfd : bool (default, True)
-            Watch the file descripttor corresponding to the replaced stream.
+            Watch the file descriptor corresponding to the replaced stream.
             This is useful if you know some underlying code will write directly
             the file descriptor by its number. It will spawn a watching thread,
             that will swap the give file descriptor for a pipe, read from the
@@ -408,19 +408,39 @@ def __init__(
 
         if (
             watchfd
-            and (sys.platform.startswith("linux") or sys.platform.startswith("darwin"))
-            and ("PYTEST_CURRENT_TEST" not in os.environ)
+            and (
+                (sys.platform.startswith("linux") or sys.platform.startswith("darwin"))
+                # Pytest set its own capture. Don't redirect from within pytest.
+                and ("PYTEST_CURRENT_TEST" not in os.environ)
+            )
+            # allow forcing watchfd (mainly for tests)
+            or watchfd == "force"
         ):
-            # Pytest set its own capture. Dont redirect from within pytest.
-
             self._should_watch = True
             self._setup_stream_redirects(name)
 
         if echo:
             if hasattr(echo, "read") and hasattr(echo, "write"):
+                # make sure we aren't trying to echo on the FD we're watching!
+                # that would cause an infinite loop, always echoing on itself
+                if self._should_watch:
+                    try:
+                        echo_fd = echo.fileno()
+                    except Exception:
+                        echo_fd = None
+
+                    if echo_fd is not None and echo_fd == self._original_stdstream_fd:
+                        # echo on the _copy_ we made during
+                        # this is the actual terminal FD now
+                        echo = io.TextIOWrapper(
+                            io.FileIO(
+                                self._original_stdstream_copy,
+                                "w",
+                            )
+                        )
                 self.echo = echo
             else:
-                msg = "echo argument must be a file like object"
+                msg = "echo argument must be a file-like object"
                 raise ValueError(msg)
 
     def isatty(self):
@@ -433,7 +453,7 @@ def isatty(self):
 
     def _setup_stream_redirects(self, name):
         pr, pw = os.pipe()
-        fno = getattr(sys, name).fileno()
+        fno = self._original_stdstream_fd = getattr(sys, name).fileno()
         self._original_stdstream_copy = os.dup(fno)
         os.dup2(pw, fno)
 
@@ -456,6 +476,10 @@ def close(self):
         if self._should_watch:
             self._should_watch = False
             self.watch_fd_thread.join()
+            # restore original FDs
+            os.dup2(self._original_stdstream_copy, self._original_stdstream_fd)
+            os.close(self._original_stdstream_copy)
+            print("closing", self)
         if self._exc:
             etype, value, tb = self._exc
             traceback.print_exception(etype, value, tb)

ipykernel/tests/test_io.py

@@ -1,6 +1,7 @@
 """Test IO capturing functionality"""
 
 import io
+import sys
 import warnings
 
 import pytest
@@ -10,20 +11,29 @@
 from ipykernel.iostream import MASTER, BackgroundSocket, IOPubThread, OutStream
 
 
-def test_io_api():
-    """Test that wrapped stdout has the same API as a normal TextIO object"""
-    session = Session()
+@pytest.fixture
+def ctx():
     ctx = zmq.Context()
-    pub = ctx.socket(zmq.PUB)
-    thread = IOPubThread(pub)
-    thread.start()
+    return ctx
+    # yield ctx
+    ctx.destroy()
 
-    stream = OutStream(session, thread, "stdout")
 
-    # cleanup unused zmq objects before we start testing
-    thread.stop()
-    thread.close()
-    ctx.term()
+@pytest.fixture
+def iopub_thread(ctx):
+    with ctx.socket(zmq.PUB) as pub:
+        thread = IOPubThread(pub)
+        thread.start()
+
+        yield thread
+        thread.stop()
+        thread.close()
+
+
+def test_io_api(iopub_thread):
+    """Test that wrapped stdout has the same API as a normal TextIO object"""
+    session = Session()
+    stream = OutStream(session, iopub_thread, "stdout")
 
     assert stream.errors is None
     assert not stream.isatty()
@@ -43,21 +53,14 @@ def test_io_api():
         stream.write(b"")  # type:ignore
 
 
-def test_io_isatty():
+def test_io_isatty(iopub_thread):
     session = Session()
-    ctx = zmq.Context()
-    pub = ctx.socket(zmq.PUB)
-    thread = IOPubThread(pub)
-    thread.start()
-
-    stream = OutStream(session, thread, "stdout", isatty=True)
+    stream = OutStream(session, iopub_thread, "stdout", isatty=True)
     assert stream.isatty()
 
 
-def test_io_thread():
-    ctx = zmq.Context()
-    pub = ctx.socket(zmq.PUB)
-    thread = IOPubThread(pub)
+def test_io_thread(iopub_thread):
+    thread = iopub_thread
     thread._setup_pipe_in()
     msg = [thread._pipe_uuid, b"a"]
     thread._handle_pipe_msg(msg)
@@ -72,40 +75,70 @@ def test_io_thread():
     thread._really_send(None)
 
 
-def test_background_socket():
-    ctx = zmq.Context()
-    pub = ctx.socket(zmq.PUB)
-    thread = IOPubThread(pub)
-    sock = BackgroundSocket(thread)
+def test_background_socket(iopub_thread):
+    sock = BackgroundSocket(iopub_thread)
     assert sock.__class__ == BackgroundSocket
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", DeprecationWarning)
         sock.linger = 101
-        assert thread.socket.linger == 101
-    assert sock.io_thread == thread
+        assert iopub_thread.socket.linger == 101
+    assert sock.io_thread == iopub_thread
     sock.send(b"hi")
 
 
-def test_outstream():
+def test_outstream(iopub_thread):
     session = Session()
-    ctx = zmq.Context()
-    pub = ctx.socket(zmq.PUB)
-    thread = IOPubThread(pub)
-    thread.start()
-
+    pub = iopub_thread.socket
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", DeprecationWarning)
         stream = OutStream(session, pub, "stdout")
-        stream = OutStream(session, thread, "stdout", pipe=object())
+        stream.close()
+        stream = OutStream(session, iopub_thread, "stdout", pipe=object())
+        stream.close()
 
-        stream = OutStream(session, thread, "stdout", watchfd=False)
+        stream = OutStream(session, iopub_thread, "stdout", watchfd=False)
         stream.close()
 
-    stream = OutStream(session, thread, "stdout", isatty=True, echo=io.StringIO())
-    with pytest.raises(io.UnsupportedOperation):
-        stream.fileno()
-    stream._watch_pipe_fd()
-    stream.flush()
-    stream.write("hi")
-    stream.writelines(["ab", "cd"])
-    assert stream.writable()
+    stream = OutStream(session, iopub_thread, "stdout", isatty=True, echo=io.StringIO())
+
+    with stream:
+        with pytest.raises(io.UnsupportedOperation):
+            stream.fileno()
+        stream._watch_pipe_fd()
+        stream.flush()
+        stream.write("hi")
+        stream.writelines(["ab", "cd"])
+        assert stream.writable()
+
+
+def test_echo_watch(capfd, iopub_thread):
+    """Test echo on underlying FD while capturing the same FD
+
+    If not careful, this
+    """
+    session = Session()
+    import os
+
+    fd_stdout = os.fdopen(sys.stdout.fileno(), "wb")
+    stream = OutStream(
+        session,
+        iopub_thread,
+        "stdout",
+        isatty=True,
+        echo=sys.stdout,
+    )
+    # fd_stdout.close()
+    save_stdout = sys.stdout
+    with stream:
+        fd_stdout.write(b"fd\n")
+        fd_stdout.flush()
+        sys.stdout = stream
+        sys.__stdout__.write("__stdout__\n")
+        sys.__stdout__.flush()
+        sys.stdout.write("stdout")
+        sys.stdout.flush()
+    sys.stdout = save_stdout
+
+    out, err = capfd.readouterr()
+    print(out, err)
+    assert out.strip() == "fd\n__stdout__\nstdout"