promote redirectIO (#74)

Co-authored-by: sseraj <[email protected]>

baseclasses/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.6.0"
+__version__ = "1.6.1"
 
 from .problems import (
     AeroProblem,

baseclasses/utils/__init__.py

@@ -1,7 +1,7 @@
 from .containers import CaseInsensitiveSet, CaseInsensitiveDict
 from .error import Error
 from .utils import getPy3SafeString, pp
-from .fileIO import writeJSON, readJSON, writePickle, readPickle
+from .fileIO import writeJSON, readJSON, writePickle, readPickle, redirectIO, redirectingIO
 
 __all__ = [
     "CaseInsensitiveSet",
@@ -13,4 +13,6 @@
     "readJSON",
     "writePickle",
     "readPickle",
+    "redirectIO",
+    "redirectingIO",
 ]

baseclasses/utils/fileIO.py

@@ -1,3 +1,7 @@
+import io
+import os
+import sys
+from contextlib import contextmanager
 import pickle
 import json
 import numpy as np
@@ -157,3 +161,100 @@ def writePickle(fname, obj, comm=None):
             pickle.dump(obj, handle)
     if comm is not None:
         comm.barrier()
+
+
+"""
+Functions for redirecting stdout/stderr to different streams
+
+Based on: http://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/.
+"""
+
+
+def redirectIO(f_out, f_err=None):
+    """
+    This function redirects stdout/stderr to the given file handle.
+
+    Parameters
+    ----------
+    f_out : file
+        A file stream to redirect stdout to
+
+    f_err : file
+        A file stream to redirect stderr to. If none is specified it is set to `f_out`
+    """
+
+    if f_err is None:
+        f_err = f_out
+
+    orig_out = sys.stdout.fileno()
+    orig_err = sys.stderr.fileno()
+
+    # flush the standard out
+    sys.stdout.flush()
+    sys.stderr.flush()
+
+    # close the standard
+    sys.stdout.close()
+    sys.stderr.close()
+
+    os.dup2(f_out.fileno(), orig_out)
+    os.dup2(f_err.fileno(), orig_err)
+
+    # reopen the stream with new file descriptors
+    sys.stdout = io.TextIOWrapper(os.fdopen(orig_out, "wb"))
+    sys.stderr = io.TextIOWrapper(os.fdopen(orig_err, "wb"))
+
+
+@contextmanager
+def redirectingIO(f_out, f_err=None):
+    """
+    A function that redirects stdout in a with block and returns to the stdout after the `with` block completes.
+    The filestream passed to this function will be closed after exiting the `with` block.
+
+    Here is an example of usage where all adflow output is redirected to the file `adflow_out.txt`:
+    >>> from baseclasses.utils import redirectIO
+    >>> print("Printing some information to terminal")
+    >>> with redirectIO.redirectingIO(open("adflow_out.txt", "w")):
+    ...     CFDSolver = ADFLOW(options=options)
+    ...     CFDSolver(AeroProblem(**apOptions)
+    >>> print("Printing some more information to terminal")
+
+    Parameters
+    ----------
+    f_out : file
+        A file stream that stdout should be redirected to
+
+    f_err : file
+        A file stream to redirect stderr to. If none is specified it is set to `f_out`
+    """
+
+    if f_err is None:
+        f_err = f_out
+
+    # save the file descriptors to restore to
+    saved_stdout_fd = os.dup(sys.stdout.fileno())
+    saved_stderr_fd = os.dup(sys.stderr.fileno())
+
+    # redirect the stdout/err streams
+    redirectIO(f_out, f_err)
+
+    # yield to the with block
+    yield
+
+    orig_out = sys.stdout.fileno()
+    orig_err = sys.stderr.fileno()
+
+    # flush output
+    sys.stderr.flush()
+    sys.stdout.flush()
+
+    # close the output
+    sys.stderr.close()
+    sys.stdout.close()
+
+    os.dup2(saved_stdout_fd, orig_out)
+    os.dup2(saved_stderr_fd, orig_err)
+
+    # reopen the standard streams with original file descriptors
+    sys.stdout = io.TextIOWrapper(os.fdopen(orig_out, "wb"))
+    sys.stderr = io.TextIOWrapper(os.fdopen(orig_err, "wb"))