tdir.py
from __future__ import print_function

import sys, os, re, tempfile, shutil

TEMP_PATH = None
SAVETY_SWITCH = True  # used in _rm(...)


def _rm(path):
    """ 
    !!! HANDLE WITH CARE !!! DANGEROUS !!! 
    Recursively remove files and directories, starting at directory <path>. 
    Returns True if one or more files couldn't be removed, else False. 
    :param path: Directory to start with. 
    :return: Error condition 
    """

    if not path:
        return True

    path = os.path.abspath(path)
    if SAVETY_SWITCH:
        if "/" not in path or len(path) <= 4:
            raise UserWarning("Savety Switch: path=%s" % path)

    try:
        shutil.rmtree(path=path)
    except shutil.Error:
        return True
    return False


class ChangeDir(object):
    """ 
    Change current directory for the time we are inside a with-statement. 
 
    To redirect print-statements to a file, use: 
        with ChangeDir(<anyDirectory>) as cd: 
            fpath = cd.print_to("test_output.txt") 
            for i in range(10): 
                print("Line %s" % i)  # output redirected to file. 
 
    Note: print_to(...) returns the complete path-/filename. 
    """

    def __init__(self, path):
        self.__current_path = os.path.abspath(path)
        self._path_onenter = None
        self._sys_stdout = None
        self._fh = None

    @property
    def current_path(self):
        return self.__current_path

    def __enter__(self):
        self._path_onenter = os.path.abspath(os.curdir)
        os.chdir(self.current_path)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if isinstance(self._fh, file):
            if not self._fh.closed:
                self._fh.close()
            self._fh = None

        if self._sys_stdout is not None:
            sys.stdout = self._sys_stdout
            self._sys_stdout = None

        os.chdir(self._path_onenter if self._path_onenter else ".")

    def print_to(self, filename):
        """Redirect following print-output to file "filename" 
        """
        filename = os.path.basename(os.path.abspath(filename))
        self._fh = open(filename, "w")
        self._sys_stdout = sys.stdout
        sys.stdout = self._fh
        return filename


class TempDirectory(ChangeDir):
    """ 
    Switch to a temporary directory were files can be created. 
 
    As long as (automatically created) TEMP_PATH is not changed, 
    the same temporary directory is used on each call. 
 
    So, in the rare case you need a garanteed clean directory, 
    set TEMP_PATH to None before using "with TempDirectory(...) as ...". 
 
    """

    def __init__(self, clean_onerr=False, clean_onexit=False):
        global TEMP_PATH
        self._clean_onerr = clean_onerr
        self._clean_onexit = clean_onexit

        if TEMP_PATH is None \
                or not os.path.isdir(TEMP_PATH) \
                or not os.access(TEMP_PATH, os.W_OK):
            # No global TEMP_PATH set or not writable.
            poser = globals().get("poser", None)
            if poser:
                # Set posers TempLocation if we run in Poser
                path = tempfile.mkdtemp(dir=poser.TempLocation())
            else:
                path = tempfile.mkdtemp()

            TEMP_PATH = path

        # TEMP_PATH should be set if we reach this point.
        super(TempDirectory, self).__init__(TEMP_PATH)

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_tb is None:
            if self._clean_onexit:
                _rm(self.current_path)
        elif self._clean_onerr:
            _rm(self.current_path)

        super(TempDirectory, self).__exit__(exc_type, exc_val, exc_tb)

    @property
    def filename(self):
        return os.path.abspath(self._fh.name)

    def remove(self, regex):
        """ 
        Remove files in TEMP_PATH by regular expression 'regex'. 
        There is no exeption generated if a file can not be removed! 
 
        :param regex: Regular expresion for files to remove. 
        :type regex: str 
        :return: False if all files could be removed, else True. 
        :rtype: bool 
        """
        err = False

        if self.current_path is not None and os.path.isdir(self.current_path):
            for p in os.listdir(self.current_path):
                if re.search(p, regex):
                    pp = os.path.join(self.current_path, p)
                    if os.path.isdir(pp):
                        err |= _rm(pp)
                    else:
                        try:
                            os.remove(pp)
                        except Exception:
                            err = True
        return err


class TempFile(TempDirectory):
    """ 
    Create a writable temporary file in directory TEMP_PATH. 
    If parameter "text" is False, file is openend "wb" (binary). 
    Parameter "clean_onerror=True" removes all files in TEMP_PATH in case of error. 
    Parameter "clean_onexit=True" removes all files in TEMP_PATH if the with-block 
    is left. 
 
    TEMP_PATH is set if it does not point to a path. 
    <TempFile>.filename returns the path to this file. 
    """

    def __init__(self, clean_onerror=True, clean_onexit=False, text=True):
        super(TempFile, self).__init__(clean_onerr=clean_onerror, clean_onexit=clean_onexit)
        self._text = text

    def __enter__(self):
        super(TempFile, self).__enter__()
        fh, self._filename = tempfile.mkstemp(dir=self.current_path, text=self._text)
        self._fh = os.fdopen(fh, "w" if self._text else "wb")
        self._sys_stdout = sys.stdout
        sys.stdout = self._fh
        return self

    @property
    def filename(self): return os.path.abspath(self._filename)

# ****************************************************************************
# ***
# *** Testing
# ***
# ****************************************************************************


def test1():
    """ 
    ======================= 
    Using class ChangeDir() 
    ======================= 
 
    Change temporarily to a certain path, write a file to it 
    and change back to the path used before. 
    Usefull if you have to deal with different files in different pathes: 
        with ChangeDir(...): 
            # handle file in this dir 
            with ChangeDir(...): 
                # handle file in that dir 
            ... back to this dir 
        ... back to inital dir 
    """
    # Create a temporary directory
    path = tempfile.mkdtemp()

    with ChangeDir(path):
        print("We are now in path '%s'." % os.path.abspath(os.curdir))
        with open("newfile.test", "w") as fh:
            print("File opened as '%s'." % os.path.abspath(fh.name))
            for i in range(10):
                print("Line", i, file=fh)

    # Automagic change back to previous directory after with.
    print("We are now back in path '%s'." % os.path.abspath(os.curdir))
    _rm(path)  # remove created temp-directory

    with ChangeDir(path) as cd:
        fname = cd.print_to("newfile.test")
        # Above line redirects all print output to file.
        for i in range(10):
            print("Line", i)  # redirected to file

    print("written to", fname)
    fh = open(fname)
    print(fh.readlines())
    fh.close()


def test2():
    """ 
    =========================== 
    Using class TempDirectory() 
    =========================== 
 
    Creates file in a temporary directory. 
 
    TEMP_PATH stays the same directory as long as TEMP_PATH isn't touched. 
    """

    # Print the name of directory we are currently in:
    currentpath = os.path.abspath(os.curdir)
    print("Current path:", currentpath)

    # Change current directory to TEMP_PATH (global var is automatically set).
    # In case of any uncatched error-exception all files and directories
    # in TEMP_PATH are removed (clean_onerror=True), so no file-fragments are left.

    with TempDirectory(clean_onerr=True):
        print("Current Temp Directory is now:", os.path.abspath(os.curdir))
        # open testfile in TEMP_PATH and write to it
        with open("test.test", "w") as fh:
            for i in range(10):
                print("Line", i, file=fh)

    # Directory is automatically restored at this point.
    print("Is restored:", currentpath == os.path.abspath(os.curdir))

    # Similar, but get a TempPath object in tp to be able to access methods:
    with TempDirectory(clean_onerr=True) as tp:
        # Remove files/directories ending with ".test" (regular expression).
        tp.remove(r"\.test$")


if __name__ == "__main__":
    # Change current directory to script path
    os.chdir(os.path.dirname(sys.argv[0]))
    test1()
    test2()
    with TempFile() as t:
        filename = t.filename
        print("dsdasdadasdasd")  # output to file
        print("......")

    print(filename)
    _rm(TEMP_PATH)