Source code for MPF.IOHelpers

try:
    from cStringIO import StringIO
except ImportError:
    from io import StringIO
import sys
import os
import threading
import tempfile
import shutil
import atexit

import ROOT

class ROpen(object):
    """ context manager to open root files """
    def __init__(self, fileName, mode='READ'):
        self.fileName = fileName
        self.mode = mode

    def __enter__(self):
        self.openFile = ROOT.TFile.Open(self.fileName, self.mode)
        if self.openFile == None:
            raise IOError("Couldn't open root file at {}".format(self.fileName))
        return self.openFile

    def __exit__(self, *args):
        self.openFile.Close()

class Capturing(list):
    """context manager to capture stdout of the encapsulated
    command(s). Might be useful to get some Information that ROOT won't
    give us, but print.
    See
    http://stackoverflow.com/questions/16571150/how-to-capture-stdout-output-from-a-python-function-call#16571630
    NOT WORKING YET SINCE IT DOESN'T CAPTURE OUTPUT FROM C++ CALLS
    """
    def __enter__(self):
        self._stdout = sys.stdout
        sys.stdout = self._stringio = StringIO()
        return self
    def __exit__(self, *args):
        self.extend(self._stringio.getvalue().splitlines())
        print(self)
        print("blub")
        del self._stringio    # free up some memory
        sys.stdout = self._stdout


class cd:
    """Context manager for changing the current working directory"""
    def __init__(self, newPath):
        self.newPath = os.path.expanduser(newPath)

    def __enter__(self):
        self.savedPath = os.getcwd()
        os.chdir(self.newPath)

    def __exit__(self, etype, value, traceback):
        os.chdir(self.savedPath)


class workInTempDir:
    """Context manager for changing to a temporary directory that will be
    deleted afterwards. The given dir will be the base directory in
    which the tempdir is created. If not given, the system tmp
    directory will be used. If *skipCleanup* is given the directory is
    not deleted afterwards. Use *prefix* to control the naming format.
    If *cleanAtExit* is set the tmp directory is cleaned at exit of the script,
    not when exiting the context.
    """
    def __init__(self, baseDir=None, skipCleanup=False, prefix="tmp_", cleanAtExit=False):
        self.baseDir = baseDir
        self.tempDir = tempfile.mkdtemp(dir=baseDir, prefix=prefix)
        self.skipCleanup = skipCleanup
        self.cleanAtExit = cleanAtExit

    def __enter__(self):
        self.savedPath = os.getcwd()
        os.chdir(self.tempDir)

    def __exit__(self, etype, value, traceback):
        os.chdir(self.savedPath)
        if self.skipCleanup:
            return
        if not self.cleanAtExit:
            self.cleanup()
        else:
            atexit.register(self.cleanup)

    def cleanup(self):
        shutil.rmtree(self.tempDir)