Source code for baseclasses.utils.fileIO

import io
import os
import sys
from contextlib import contextmanager
import pickle
import json
import numpy as np
from .containers import CaseInsensitiveDict, CaseInsensitiveSet


[docs] def writeJSON(fname, obj, comm=None): """ Write an object to a JSON file. This includes a custom NumPy encoder to reliably write NumPy arrays to JSON, which can then be read back via :meth:`readJSON`. Parameters ---------- fname : str The file name obj : dict or ndarray The object to be written to JSON comm : mpi4py.MPI.Comm, optional The communicator over which this function is called. If supplied, only the root proc will be used for file IO. """ class MyEncoder(json.JSONEncoder): """ Custom encoder class for Numpy arrays and CaseInsensitiveDict """ def default(self, o): """ If input object is an ndarray it will be converted into a dict holding dtype, shape and the data, base64 encoded. """ if isinstance(o, np.ndarray): if o.flags["C_CONTIGUOUS"]: pass else: o = np.ascontiguousarray(o) assert o.flags["C_CONTIGUOUS"] if o.size == 1: return o.item() else: return dict(__ndarray__=o.tolist(), dtype=str(o.dtype), shape=o.shape) elif isinstance(o, np.integer): return dict(__ndarray__=int(o), dtype=str(o.dtype), shape=o.shape) elif isinstance(o, np.floating): return dict(__ndarray__=float(o), dtype=str(o.dtype), shape=o.shape) elif isinstance(o, CaseInsensitiveDict): return dict(o) elif isinstance(o, CaseInsensitiveSet): return set(o) else: # Let the base class default method raise the TypeError super().default(o) if (comm is None) or (comm is not None and comm.rank == 0): with open(fname, "w") as json_file: json.dump(obj, json_file, sort_keys=True, indent=4, separators=(",", ": "), cls=MyEncoder) if comm is not None: comm.barrier()
[docs] def readJSON(fname, comm=None): """ Reads a JSON file and return the contents as a dictionary. This includes a custom NumPy reader to retrieve NumPy arrays, matching the :meth:`writeJSON` function. Parameters ---------- file_name : str The file name comm : mpi4py.MPI.Comm, optional The communicator over which this function is called. If supplied, only the root proc will be used for file IO. References ---------- This is based on `this stack overflow answer <https://stackoverflow.com/questions/3488934/simplejson-and-numpy-array/24375113#24375113>`_ """ def json_numpy_obj_hook(dct): """ Decodes a previously encoded numpy ndarray with proper shape and dtype. Parameters ---------- dct: dictionary JSON dictionary containing the encoded ndarray Returns ------- dct: dictionary or Numpy array The decoded Numpy array, or the input dct if it is not an encoded Numpy array """ if isinstance(dct, dict) and "__ndarray__" in dct: data = dct["__ndarray__"] return np.array(data, dct["dtype"]).reshape(dct["shape"]) return dct if not os.path.isfile(fname): raise FileNotFoundError(f"The JSON file {fname} cannot be found.") data = None if (comm is None) or (comm is not None and comm.rank == 0): with open(fname, "r") as json_file: data = json.load(json_file, object_hook=json_numpy_obj_hook) if comm is not None: data = comm.bcast(data) return data
[docs] def readPickle(fname, comm=None): """ This is a parallel pickle.load function, which is performed on the root proc only. Error checking is necessary to provide py2 compatibility. Parameters ---------- fname : str The pickle file name comm : mpi4py.MPI.Comm, optional The communicator over which this function is called. If supplied, only the root proc will be used for file IO. Returns ------- obj : The object stored in the pickle file """ if not os.path.isfile(fname): raise FileNotFoundError(f"The pickle file {fname} cannot be found.") obj = None if (comm is None) or (comm is not None and comm.rank == 0): try: with open(fname, "rb") as f: obj = pickle.load(f) except UnicodeDecodeError: # if pickled with py2 with open(fname, "rb") as f: obj = pickle.load(f, encoding="latin1") if comm is not None: comm.barrier() obj = comm.bcast(obj) return obj
[docs] def writePickle(fname, obj, comm=None): """ Parallel pickle writing function, only performs operations on the root proc Parameters ---------- fname : str The pickle file name obj : any object which can be pickled by Python The object to be pickled comm : mpi4py.MPI.Comm, optional The communicator over which this function is called. If supplied, only the root proc will be used for file IO. """ if (comm is None) or (comm is not None and comm.rank == 0): with open(fname, "wb") as handle: 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/. """
[docs] 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"))
[docs] @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) # close copies os.close(saved_stdout_fd) os.close(saved_stderr_fd) # 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"))