mpl CatchOutput bugs

mpl_catchoutput.py

#!/usr/bin/env python
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import matplotlib
matplotlib.use("agg")

from contextlib import contextmanager
import tempfile
import sys
import os
import numpy as np

@contextmanager
def CatchOutput():
    """
    A context manager that catches stdout/stderr/exit() for its scope.

    Always use with "with" statement. Does nothing otherwise.

    Roughly based on: http://stackoverflow.com/a/17954769

    This variant does not leak file descriptors.

    >>> with CatchOutput() as out:  # doctest: +SKIP
    ...    os.system('echo "mystdout"')
    ...    os.system('echo "mystderr" >&2')
    >>> print(out.stdout)  # doctest: +SKIP
    mystdout
    >>> print(out.stderr)  # doctest: +SKIP
    mystderr
    """
    stdout_file, stdout_filename = tempfile.mkstemp(prefix="obspy-")
    stderr_file, stderr_filename = tempfile.mkstemp(prefix="obspy-")

    try:
        fd_stdout = sys.stdout.fileno()
        fd_stderr = sys.stderr.fileno()

        # Dummy class to transport the output.
        class Output():
            pass
        out = Output()
        out.stdout = ""
        out.stderr = ""

        with os.fdopen(os.dup(sys.stdout.fileno()), "wb") as old_stdout:
            with os.fdopen(os.dup(sys.stderr.fileno()), "wb") as old_stderr:
                sys.stdout.flush()
                sys.stderr.flush()

                os.dup2(stdout_file, fd_stdout)
                os.dup2(stderr_file, fd_stderr)

                os.close(stdout_file)
                os.close(stderr_file)

                try:
                    raised = False
                    yield out
                except SystemExit:
                    raised = True
                finally:
                    sys.stdout.flush()
                    sys.stderr.flush()
                    os.fsync(sys.stdout.fileno())
                    os.fsync(sys.stderr.fileno())
                    sys.stdout = sys.__stdout__
                    sys.stderr = sys.__stderr__
                    sys.stdout.flush()
                    sys.stderr.flush()
                    os.dup2(old_stdout.fileno(), sys.stdout.fileno())
                    os.dup2(old_stderr.fileno(), sys.stderr.fileno())

                    with open(stdout_filename, "rb") as fh:
                        out.stdout = fh.read()
                    with open(stderr_filename, "rb") as fh:
                        out.stderr = fh.read()

                    if raised:
                        raise SystemExit(out.stderr)

    finally:
        # Make sure to always close and remove the temporary files.
        try:
            os.close(stdout_file)
        except:
            pass
        try:
            os.close(stderr_file)
        except:
            pass
        try:
            os.remove(stdout_filename)
        except OSError:
            pass
        try:
            os.remove(stderr_filename)
        except OSError:
            pass


def main(output):
    import matplotlib.pyplot as plt

    fig = plt.figure()
    ax = fig.add_subplot(111)

    startend = np.array([[733055, 733057], [733056, 733058]])

    ax.plot(startend[:, 0], startend[:, 1])

    fig.savefig(output, dpi=50)


if 0:
    with CatchOutput():
        main('x1.png')
        main('x2.png')
else:
    with CatchOutput():
        main('x1.png')
    with CatchOutput():
        main('x2.png')