diff --git a/src/calibre/ptempfile.py b/src/calibre/ptempfile.py index 60188c5631..ad8da9dcb5 100644 --- a/src/calibre/ptempfile.py +++ b/src/calibre/ptempfile.py @@ -4,54 +4,15 @@ __copyright__ = '2008, Kovid Goyal ' Provides platform independent temporary files that persist even after being closed. ''' -import atexit import os import tempfile from calibre.constants import __appname__, __version__, filesystem_encoding, get_windows_temp_path, ismacos, iswindows - - -def cleanup(path): - try: - import os as oss - if oss.path.exists(path): - oss.remove(path) - except: - pass - +from calibre.utils.safe_atexit import remove_dir, remove_file_atexit, remove_folder_atexit, unlink _base_dir = None -def remove_dir(x): - try: - import shutil - shutil.rmtree(x, ignore_errors=True) - except: - pass - - -def determined_remove_dir(x): - for i in range(10): - try: - import shutil - shutil.rmtree(x) - return - except: - import os # noqa - if os.path.exists(x): - # In case some other program has one of the temp files open. - import time - time.sleep(0.1) - else: - return - try: - import shutil - shutil.rmtree(x, ignore_errors=True) - except: - pass - - def app_prefix(prefix): if iswindows: return f'{__appname__}_' @@ -81,6 +42,9 @@ def osx_cache_dir(): return q +get_default_tempdir = tempfile.gettempdir + + def base_dir(): global _base_dir if _base_dir is not None and not os.path.exists(_base_dir): @@ -119,23 +83,25 @@ def base_dir(): # https://developer.apple.com/library/mac/documentation/Darwin/Reference/ManPages/man3/confstr.3.html base = osx_cache_dir() - _base_dir = tempfile.mkdtemp(prefix=prefix, dir=base) - atexit.register(determined_remove_dir if iswindows else remove_dir, _base_dir) - - try: - tempfile.gettempdir() - except Exception: - # Widows temp vars set to a path not encodable in mbcs - # Use our temp dir - tempfile.tempdir = _base_dir + _base_dir = tempfile.mkdtemp(prefix=prefix, dir=base or get_default_tempdir()) + remove_folder_atexit(_base_dir) return _base_dir +def fix_tempfile_module(): + # We want the tempfile module to use base_dir() as its tempdir, but we dont + # want to call base_dir() now as it will possibly create a tempdir, do that + # only on demand. + global get_default_tempdir + if tempfile._gettempdir is not base_dir: + get_default_tempdir = tempfile._gettempdir + tempfile._gettempdir = base_dir + + def reset_base_dir(): global _base_dir _base_dir = None - base_dir() def force_unicode(x): @@ -173,7 +139,7 @@ class PersistentTemporaryFile: self._file = os.fdopen(fd, mode) self._name = name self._fd = fd - atexit.register(cleanup, name) + remove_file_atexit(name) def __getattr__(self, name): if name == 'name': @@ -201,8 +167,7 @@ def PersistentTemporaryDirectory(suffix='', prefix='', dir=None): if dir is None: dir = base_dir() tdir = _make_dir(suffix, prefix, dir) - - atexit.register(remove_dir, tdir) + remove_folder_atexit(tdir) return tdir @@ -249,7 +214,7 @@ class TemporaryFile: return name def __exit__(self, *args): - cleanup(self._name) + unlink(self._name) class SpooledTemporaryFile(tempfile.SpooledTemporaryFile): diff --git a/src/calibre/startup.py b/src/calibre/startup.py index 5d510f1efe..01f364a2a1 100644 --- a/src/calibre/startup.py +++ b/src/calibre/startup.py @@ -72,15 +72,10 @@ def initialize_calibre(): if hasattr(initialize_calibre, 'initialized'): return initialize_calibre.initialized = True - # Ensure that all temp files/dirs are created under a calibre tmp dir - from calibre.ptempfile import base_dir - try: - base_dir() - except OSError: - pass # Ignore this error during startup, so we can show a better error message to the user later. + from calibre.ptempfile import fix_tempfile_module + fix_tempfile_module() - # # Ensure that the max number of open files is at least 1024 if iswindows: # See https://msdn.microsoft.com/en-us/library/6e3b887c.aspx diff --git a/src/calibre/utils/safe_atexit.py b/src/calibre/utils/safe_atexit.py index 45348b78f1..35d7a636c3 100644 --- a/src/calibre/utils/safe_atexit.py +++ b/src/calibre/utils/safe_atexit.py @@ -104,7 +104,14 @@ def main(): traceback.print_exc() -def main_for_test(do_forced_exit=False): +def main_for_test(do_forced_exit=False, check_tdir=False): + if check_tdir: + import tempfile + + from calibre.ptempfile import base_dir + print(tempfile.gettempdir()) + print(base_dir()) + return tf = 'test-folder' os.mkdir(tf) open(os.path.join(tf, 'test-file'), 'w').close() @@ -146,5 +153,9 @@ def find_tests(): p = start_pipe_worker('from calibre.utils.safe_atexit import main_for_test; main_for_test(True)', cwd=tdir) p.wait(10) self.wait_for_empty(tdir) + p = start_pipe_worker('from calibre.utils.safe_atexit import main_for_test; main_for_test(check_tdir=True)') + tempfiledir, bdir = p.stdout.read().decode().splitlines()[:2] + p.wait(10) + self.assertEqual(bdir, tempfiledir) return unittest.defaultTestLoader.loadTestsFromTestCase(TestSafeAtexit)