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)