Issue 5178: Add tempfile.TemporaryDirectory (original patch by Neil Schemenauer)

5 files changed, 229 insertions, 3 deletions

diff --git a/Doc/library/tempfile.rst b/Doc/library/tempfile.rst
index a13df0dd22..5caea677e3 100644
--- a/Doc/library/tempfile.rst
+++ b/Doc/library/tempfile.rst
@@ -25,7 +25,7 @@ no longer necessary to use the global *tempdir* and *template* variables.
 To maintain backward compatibility, the argument order is somewhat odd; it
 is recommended to use keyword arguments for clarity.
 
-The module defines the following user-callable functions:
+The module defines the following user-callable items:
 
 .. function:: TemporaryFile(mode='w+b', buffering=None, encoding=None, newline=None, suffix='', prefix='tmp', dir=None)
 
@@ -83,6 +83,24 @@ The module defines the following user-callable functions:
    used in a :keyword:`with` statement, just like a normal file.
 
 
+.. function:: TemporaryDirectory(suffix='', prefix='tmp', dir=None)
+
+   This function creates a temporary directory using :func:`mkdtemp`
+   (the supplied arguments are passed directly to the underlying function).
+   The resulting object can be used as a context manager (see
+   :ref:`context-managers`).  On completion of the context (or destruction
+   of the temporary directory object), the newly created temporary directory
+   and all its contents are removed from the filesystem.
+
+   The directory name can be retrieved from the :attr:`name` member
+   of the returned object.
+
+   The directory can be explicitly cleaned up by calling the
+   :func:`cleanup` method.
+
+   .. versionadded:: 3.2
+
+
 .. function:: mkstemp(suffix='', prefix='tmp', dir=None, text=False)
 
    Creates a temporary file in the most secure manner possible.  There are
@@ -210,3 +228,36 @@ the appropriate function arguments, instead.
    Return the filename prefix used to create temporary files.  This does not
    contain the directory component.
 
+
+Examples
+--------
+
+Here are some examples of typical usage of the :mod:`tempfile` module::
+
+    >>> import tempfile
+
+    # create a temporary file and write some data to it
+    >>> fp = tempfile.TemporaryFile()
+    >>> fp.write('Hello world!')
+    # read data from file
+    >>> fp.seek(0)
+    >>> fp.read()
+    'Hello world!'
+    # close the file, it will be removed
+    >>> fp.close()
+
+    # create a temporary file using a context manager
+    >>> with tempfile.TemporaryFile() as fp:
+    ...     fp.write('Hello world!')
+    ...     fp.seek(0)
+    ...     fp.read()
+    'Hello world!'
+    >>>
+    # file is now closed and removed
+
+    # create a temporary directory using the context manager
+    >>> with tempfile.TemporaryDirectory() as tmpdirname:
+    ...     print 'created temporary directory', tmpdirname
+    >>>
+    # directory and contents have been removed
+
diff --git a/Doc/whatsnew/3.2.rst b/Doc/whatsnew/3.2.rst
index c16fe8741a..386e527af8 100644
--- a/Doc/whatsnew/3.2.rst
+++ b/Doc/whatsnew/3.2.rst
@@ -496,6 +496,13 @@ New, Improved, and Deprecated Modules
 
   (Contributed by Giampaolo Rodolà; :issue:`6706`.)
 
+* The :mod:`tempfile` module has a new context manager,
+  :class:`~tempfile.TemporaryDirectory` which provides easy deterministic
+  cleanup of temporary directories.
+
+  (Contributed by Neil Schemenauer and Nick Coghlan; :issue:`5178`.)
+
+
 Multi-threading
 ===============
 
diff --git a/Lib/tempfile.py b/Lib/tempfile.py
index 049cdaa2c2..699fd0a9f9 100644
--- a/Lib/tempfile.py
+++ b/Lib/tempfile.py
@@ -19,7 +19,7 @@ This module also provides some data items to the user:
 
 __all__ = [
     "NamedTemporaryFile", "TemporaryFile", # high level safe interfaces
-    "SpooledTemporaryFile",
+    "SpooledTemporaryFile", "TemporaryDirectory",
     "mkstemp", "mkdtemp",                  # low level safe interfaces
     "mktemp",                              # deprecated unsafe interface
     "TMP_MAX", "gettempprefix",            # constants
@@ -613,3 +613,66 @@ class SpooledTemporaryFile:
 
     def xreadlines(self, *args):
         return self._file.xreadlines(*args)
+
+
+class TemporaryDirectory(object):
+    """Create and return a temporary directory.  This has the same
+    behavior as mkdtemp but can be used as a context manager.  For
+    example:
+
+        with TemporaryDirectory() as tmpdir:
+            ...
+
+    Upon exiting the context, the directory and everthing contained
+    in it are removed.
+    """
+
+    def __init__(self, suffix="", prefix=template, dir=None):
+        self.name = mkdtemp(suffix, prefix, dir)
+        self._closed = False
+
+    def __enter__(self):
+        return self.name
+
+    def cleanup(self):
+        if not self._closed:
+            self._rmtree(self.name)
+            self._closed = True
+
+    def __exit__(self, exc, value, tb):
+        self.cleanup()
+
+    __del__ = cleanup
+
+
+    # XXX (ncoghlan): The following code attempts to make
+    # this class tolerant of the module nulling out process
+    # that happens during CPython interpreter shutdown
+    # Alas, it doesn't actually manage it. See issue #10188
+    _listdir = staticmethod(_os.listdir)
+    _path_join = staticmethod(_os.path.join)
+    _isdir = staticmethod(_os.path.isdir)
+    _remove = staticmethod(_os.remove)
+    _rmdir = staticmethod(_os.rmdir)
+    _os_error = _os.error
+
+    def _rmtree(self, path):
+        # Essentially a stripped down version of shutil.rmtree.  We can't
+        # use globals because they may be None'ed out at shutdown.
+        for name in self._listdir(path):
+            fullname = self._path_join(path, name)
+            try:
+                isdir = self._isdir(fullname)
+            except self._os_error:
+                isdir = False
+            if isdir:
+                self._rmtree(fullname)
+            else:
+                try:
+                    self._remove(fullname)
+                except self._os_error:
+                    pass
+        try:
+            self._rmdir(path)
+        except self._os_error:
+            pass
diff --git a/Lib/test/test_tempfile.py b/Lib/test/test_tempfile.py
index b0976d225a..03e9a46274 100644
--- a/Lib/test/test_tempfile.py
+++ b/Lib/test/test_tempfile.py
@@ -85,7 +85,8 @@ class test_exports(TC):
             "gettempdir" : 1,
             "tempdir" : 1,
             "template" : 1,
-            "SpooledTemporaryFile" : 1
+            "SpooledTemporaryFile" : 1,
+            "TemporaryDirectory" : 1,
         }
 
         unexp = []
@@ -889,6 +890,107 @@ class test_TemporaryFile(TC):
 if tempfile.NamedTemporaryFile is not tempfile.TemporaryFile:
     test_classes.append(test_TemporaryFile)
 
+
+# Helper for test_del_on_shutdown
+class NulledModules:
+    def __init__(self, *modules):
+        self.refs = [mod.__dict__ for mod in modules]
+        self.contents = [ref.copy() for ref in self.refs]
+
+    def __enter__(self):
+        for d in self.refs:
+            for key in d:
+                d[key] = None
+
+    def __exit__(self, *exc_info):
+        for d, c in zip(self.refs, self.contents):
+            d.clear()
+            d.update(c)
+
+class test_TemporaryDirectory(TC):
+    """Test TemporaryDirectory()."""
+
+    def do_create(self, dir=None, pre="", suf="", recurse=1):
+        if dir is None:
+            dir = tempfile.gettempdir()
+        try:
+            tmp = tempfile.TemporaryDirectory(dir=dir, prefix=pre, suffix=suf)
+        except:
+            self.failOnException("TemporaryDirectory")
+        self.nameCheck(tmp.name, dir, pre, suf)
+        # Create a subdirectory and some files
+        if recurse:
+            self.do_create(tmp.name, pre, suf, recurse-1)
+        with open(os.path.join(tmp.name, "test.txt"), "wb") as f:
+            f.write(b"Hello world!")
+        return tmp
+
+    def test_explicit_cleanup(self):
+        # A TemporaryDirectory is deleted when cleaned up
+        dir = tempfile.mkdtemp()
+        try:
+            d = self.do_create(dir=dir)
+            self.assertTrue(os.path.exists(d.name),
+                            "TemporaryDirectory %s does not exist" % d.name)
+            d.cleanup()
+            self.assertFalse(os.path.exists(d.name),
+                        "TemporaryDirectory %s exists after cleanup" % d.name)
+        finally:
+            os.rmdir(dir)
+
+    @support.cpython_only
+    def test_del_on_collection(self):
+        # A TemporaryDirectory is deleted when garbage collected
+        dir = tempfile.mkdtemp()
+        try:
+            d = self.do_create(dir=dir)
+            name = d.name
+            del d # Rely on refcounting to invoke __del__
+            self.assertFalse(os.path.exists(name),
+                        "TemporaryDirectory %s exists after __del__" % name)
+        finally:
+            os.rmdir(dir)
+
+    @unittest.expectedFailure # See issue #10188
+    def test_del_on_shutdown(self):
+        # A TemporaryDirectory may be cleaned up during shutdown
+        # Make sure it works with the relevant modules nulled out
+        dir = tempfile.mkdtemp()
+        try:
+            d = self.do_create(dir=dir)
+            # Mimic the nulling out of modules that
+            # occurs during system shutdown
+            modules = [os, os.path]
+            if has_stat:
+                modules.append(stat)
+            with NulledModules(*modules):
+                d.cleanup()
+            self.assertFalse(os.path.exists(d.name),
+                        "TemporaryDirectory %s exists after cleanup" % d.name)
+        finally:
+            os.rmdir(dir)
+
+    def test_multiple_close(self):
+        # Can be cleaned-up many times without error
+        d = self.do_create()
+        d.cleanup()
+        try:
+            d.cleanup()
+            d.cleanup()
+        except:
+            self.failOnException("cleanup")
+
+    def test_context_manager(self):
+        # Can be used as a context manager
+        d = self.do_create()
+        with d as name:
+            self.assertTrue(os.path.exists(name))
+            self.assertEqual(name, d.name)
+        self.assertFalse(os.path.exists(name))
+
+
+test_classes.append(test_TemporaryDirectory)
+
 def test_main():
     support.run_unittest(*test_classes)
 
diff --git a/Misc/NEWS b/Misc/NEWS
index 98bc24f1cc..e1b818d36d 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -51,6 +51,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #5178: Added tempfile.TemporaryDirectory class that can be used
+  as a context manager.
+
 - Issue #1349106: Generator (and BytesGenerator) flatten method and Header
   encode method now support a 'linesep' argument.