Merge pull request #3497 from jsiirola/capture_output_jupyter

Resolve errors in `TeeStream` and `capture_output`

Author: blnicho

pyomo/common/tee.py

@@ -194,10 +194,12 @@ class capture_output(object):
 
     capture_fd : bool
 
-        If True, we will also redirect the low-level file descriptors
-        associated with stdout (1) and stderr (2) to the ``output``.
-        This is useful for capturing output emitted directly to the
-        process stdout / stderr by external compiled modules.
+        If True, we will also redirect the process file descriptors
+        ``1`` (stdout), ``2`` (stderr), and the file descriptors from
+        ``sys.stdout.fileno()`` and ``sys.stderr.fileno()`` to the
+        ``output``.  This is useful for capturing output emitted
+        directly to the process stdout / stderr by external compiled
+        modules.
 
     Returns
     -------
@@ -231,19 +233,59 @@ def __enter__(self):
         sys.stdout = self.tee.STDOUT
         sys.stderr = self.tee.STDERR
         if self.capture_fd:
-            self.fd_redirect = (
-                redirect_fd(1, self.tee.STDOUT.fileno(), synchronize=False),
-                redirect_fd(2, self.tee.STDERR.fileno(), synchronize=False),
+            tee_fd = (self.tee.STDOUT.fileno(), self.tee.STDERR.fileno())
+            self.fd_redirect = []
+            for i in range(2):
+                # Redirect the standard process file descriptor (1 or 2)
+                self.fd_redirect.append(
+                    redirect_fd(i + 1, tee_fd[i], synchronize=False)
+                )
+                # Redirect the file descriptor currently associated with
+                # sys.stdout / sys.stderr
+                try:
+                    fd = self.old[i].fileno()
+                except (AttributeError, OSError):
+                    pass
+                else:
+                    if fd != i + 1:
+                        self.fd_redirect.append(
+                            redirect_fd(fd, tee_fd[i], synchronize=False)
+                        )
+            for fdr in self.fd_redirect:
+                fdr.__enter__()
+        # We have an issue where we are (very aggressively)
+        # commandeering the terminal.  This is what we intend, but the
+        # side effect is that any errors generated by this module (e.g.,
+        # because the user gave us an invalid output stream) get
+        # completely suppressed.  So, we will make an exception to the
+        # output that we are catching and let messages logged to THIS
+        # logger to still be emitted.
+        if self.capture_fd:
+            # Because we are also comandeering the FD that underlies
+            # self.old[1], we cannot just write to that stream and
+            # instead open a new stream to the original FD.
+            #
+            # Note that we need to duplicate the FD from the redirector,
+            # as it will close the (temporary) `original_fd` descriptor
+            # when it restores the actual original descriptor
+            self.temp_log_stream = os.fdopen(
+                os.dup(self.fd_redirect[-1].original_fd), mode="w", closefd=True
             )
-            self.fd_redirect[0].__enter__()
-            self.fd_redirect[1].__enter__()
+        else:
+            self.temp_log_stream = self.old[1]
+        self.temp_log_handler = logging.StreamHandler(self.temp_log_stream)
+        logger.addHandler(self.temp_log_handler)
+        self._propagate = logger.propagate
+        logger.propagate = False
         return self.output_stream
 
     def __exit__(self, et, ev, tb):
+        # Restore any file descriptors we comandeered
         if self.fd_redirect is not None:
-            self.fd_redirect[1].__exit__(et, ev, tb)
-            self.fd_redirect[0].__exit__(et, ev, tb)
+            for fdr in reversed(self.fd_redirect):
+                fdr.__exit__(et, ev, tb)
             self.fd_redirect = None
+        # Check and restore sys.stderr / sys.stdout
         FAIL = self.tee.STDOUT is not sys.stdout
         self.tee.__exit__(et, ev, tb)
         if self.output_stream is not self.output:
@@ -252,6 +294,15 @@ def __exit__(self, et, ev, tb):
         self.old = None
         self.tee = None
         self.output_stream = None
+        # Clean up our temporary override of the local logger
+        self.temp_log_handler.flush()
+        logger.removeHandler(self.temp_log_handler)
+        if self.capture_fd:
+            self.temp_log_stream.flush()
+            self.temp_log_stream.close()
+        logger.propagate = self._propagate
+        self.temp_log_stream = None
+        self.temp_log_handler = None
         if FAIL:
             raise RuntimeError('Captured output does not match sys.stdout.')
 
@@ -378,47 +429,66 @@ def writeOutputBuffer(self, ostreams, flush):
         if not ostring:
             return
 
-        for local_stream, user_stream in ostreams:
+        for stream in ostreams:
             try:
-                written = local_stream.write(ostring)
+                written = stream.write(ostring)
             except:
-                written = 0
+                my_repr = "<%s.%s @ %s>" % (
+                    stream.__class__.__module__,
+                    stream.__class__.__name__,
+                    hex(id(stream)),
+                )
+                if my_repr in ostring:
+                    # In the case of nested capture_outputs, all the
+                    # handlers are left on the logger.  We want to make
+                    # sure that we don't create an infinite loop by
+                    # re-printing a message *this* object generated.
+                    continue
+                et, e, tb = sys.exc_info()
+                msg = "Error writing to output stream %s:\n    %s: %s\n" % (
+                    my_repr,
+                    et.__name__,
+                    e,
+                )
+                if getattr(stream, 'closed', False):
+                    msg += "Output stream closed before all output was written to it."
+                else:
+                    msg += "Is this a writeable TextIOBase object?"
+                logger.error(
+                    f"{msg}\nThe following was left in the output buffer:\n"
+                    f"    {ostring!r}"
+                )
+                continue
             if flush or (written and not self.buffering):
-                local_stream.flush()
-                if local_stream is not user_stream:
-                    user_stream.flush()
+                stream.flush()
             # Note: some derived file-like objects fail to return the
             # number of characters written (and implicitly return None).
             # If we get None, we will just assume that everything was
             # fine (as opposed to tossing the incomplete write error).
             if written is not None and written != len(ostring):
+                my_repr = "<%s.%s @ %s>" % (
+                    stream.__class__.__module__,
+                    stream.__class__.__name__,
+                    hex(id(stream)),
+                )
+                if my_repr in ostring[written:]:
+                    continue
                 logger.error(
-                    "Output stream (%s) closed before all output was "
-                    "written to it. The following was left in "
-                    "the output buffer:\n\t%r" % (local_stream, ostring[written:])
+                    "Incomplete write to output stream %s.\nThe following was "
+                    "left in the output buffer:\n    %r" % (my_repr, ostring[written:])
                 )
 
 
 class TeeStream(object):
     def __init__(self, *ostreams, encoding=None, buffering=-1):
-        self.ostreams = []
+        self.ostreams = ostreams
         self.encoding = encoding
         self.buffering = buffering
         self._stdout = None
         self._stderr = None
         self._handles = []
         self._active_handles = []
         self._threads = []
-        for user_stream in ostreams:
-            try:
-                fileno = user_stream.fileno()
-            except:
-                self.ostreams.append((user_stream, user_stream))
-                continue
-            local_stream = os.fdopen(
-                os.dup(fileno), mode=getattr(user_stream, 'mode', None), closefd=True
-            )
-            self.ostreams.append((local_stream, user_stream))
 
     @property
     def STDOUT(self):
@@ -499,9 +569,6 @@ def close(self, in_exception=False):
         self._active_handles.clear()
         self._stdout = None
         self._stderr = None
-        for local, orig in self.ostreams:
-            if orig is not local:
-                local.close()
 
     def __enter__(self):
         return self

pyomo/common/tests/test_tee.py

@@ -12,6 +12,7 @@
 
 import gc
 import itertools
+import logging
 import os
 import platform
 import time
@@ -35,7 +36,10 @@ def __init__(self):
 
     def write(self, data):
         for line in data.splitlines():
-            self.buf.append((time.time(), float(line.strip())))
+            line = line.strip()
+            if not line:
+                continue
+            self.buf.append((time.time(), float(line)))
 
     def writelines(self, data):
         for line in data:
@@ -216,10 +220,43 @@ def test_decoder_and_buffer_errors(self):
             with tee.TeeStream(out) as t:
                 out.close()
                 t.STDOUT.write("hi\n")
+        _id = hex(id(out))
         self.assertRegex(
             log.getvalue(),
-            r"^Output stream \(<.*?>\) closed before all output was written "
-            r"to it. The following was left in the output buffer:\n\t'hi\\n'\n$",
+            f"Error writing to output stream <_?io.StringIO @ {_id}>:"
+            r"\n.*\nOutput stream closed before all output was written to it.\n"
+            r"The following was left in the output buffer:\n    'hi\\n'\n$",
+        )
+
+        # TeeStream expects stream-like objects
+        out = logging.getLogger()
+        log = StringIO()
+        with LoggingIntercept(log):
+            with tee.TeeStream(out) as t:
+                t.STDOUT.write("hi\n")
+        _id = hex(id(out))
+        self.assertRegex(
+            log.getvalue(),
+            f"Error writing to output stream <logging.RootLogger @ {_id}>:"
+            r"\n.*\nIs this a writeable TextIOBase object\?\n"
+            r"The following was left in the output buffer:\n    'hi\\n'\n$",
+        )
+
+        # Catch partial writes
+        class fake_stream:
+            def write(self, data):
+                return 1
+
+        out = fake_stream()
+        log = StringIO()
+        with LoggingIntercept(log):
+            with tee.TeeStream(out) as t:
+                t.STDOUT.write("hi\n")
+        _id = hex(id(out))
+        self.assertRegex(
+            log.getvalue(),
+            f"Incomplete write to output stream <.*fake_stream @ {_id}>."
+            r"\nThe following was left in the output buffer:\n    'i\\n'\n$",
         )
 
     def test_capture_output(self):
@@ -267,6 +304,32 @@ def test_capture_output_stack_error(self):
         finally:
             sys.stdout, sys.stderr = old
 
+    def test_capture_output_invalid_ostream(self):
+        # Test that capture_output does not suppress errors from the tee
+        # module
+        _id = hex(id(15))
+        with tee.capture_output(capture_fd=True) as OUT:
+            with tee.capture_output(15):
+                sys.stderr.write("hi\n")
+        self.assertEqual(
+            OUT.getvalue(),
+            f"Error writing to output stream <builtins.int @ {_id}>:\n"
+            "    AttributeError: 'int' object has no attribute 'write'\n"
+            "Is this a writeable TextIOBase object?\n"
+            "The following was left in the output buffer:\n    'hi\\n'\n",
+        )
+
+        with tee.capture_output(capture_fd=True) as OUT:
+            with tee.capture_output(15, capture_fd=True):
+                print("hi")
+        self.assertEqual(
+            OUT.getvalue(),
+            f"Error writing to output stream <builtins.int @ {_id}>:\n"
+            "    AttributeError: 'int' object has no attribute 'write'\n"
+            "Is this a writeable TextIOBase object?\n"
+            "The following was left in the output buffer:\n    'hi\\n'\n",
+        )
+
     def test_deadlock(self):
         class MockStream(object):
             def write(self, data):
@@ -323,14 +386,15 @@ def test_buffered_stdout(self):
             sys.stdout.write(f"{time.time()}\n")
             time.sleep(self.dt)
         ts.write(f"{time.time()}\n")
-        baseline = [[(0, 0), (1, 0), (1, 0), (1, 1)]]
-        if fd:
-            # TODO: [JDS] If we are capturing the file descriptor, the
-            # stdout channel is sometimes no longer buffered.  I am not
-            # exactly sure why (my guess is because the underlying pipe
-            # is not buffered), but as it is generally not a problem to
-            # not buffer, we will put off "fixing" it.
-            baseline.append([(0, 0), (0, 0), (0, 0), (1, 1)])
+        baseline = [
+            [(0, 0), (1, 0), (1, 0), (1, 1)],
+            # TODO: [JDS] The stdout channel appears to sometimes be no
+            # longer buffered.  I am not exactly sure why (my guess is
+            # because the underlying pipe is not buffered), but as it is
+            # generally not a problem to not buffer, we will put off
+            # "fixing" it.
+            [(0, 0), (0, 0), (0, 0), (1, 1)],
+        ]
         if not ts.check(*baseline):
             self.fail(ts.error)