1# coding: utf-8
2from __future__ import absolute_import, division, print_function
3
4import os as _os
5import sys as _sys
6import warnings as _warnings
7from tempfile import mkdtemp
8
9
10class TemporaryDirectory(object):
11    """Create and return a temporary directory.  This has the same
12    behavior as mkdtemp but can be used as a context manager.  For
13    example:
14
15        with TemporaryDirectory() as tmpdir:
16            ...
17
18    Upon exiting the context, the directory and everything contained
19    in it are removed.
20    """
21
22    def __init__(self, suffix="", prefix="tmp", dir=None):
23        self._closed = False
24        self.name = None  # Handle mkdtemp raising an exception
25        self.name = mkdtemp(suffix, prefix, dir)
26
27    def __repr__(self):
28        return "<{} {!r}>".format(self.__class__.__name__, self.name)
29
30    def __enter__(self):
31        return self.name
32
33    def cleanup(self):
34        if self.name and not self._closed:
35            try:
36                self._rmtree(self.name)
37            except (TypeError, AttributeError) as ex:
38                # Issue #10188: Emit a warning on stderr
39                # if the directory could not be cleaned
40                # up due to missing globals
41                if "None" not in str(ex):
42                    raise
43                print(
44                    "ERROR: {!r} while cleaning up {!r}".format(ex, self),
45                    file=_sys.stderr,
46                )
47                return
48            self._closed = True
49
50    def __exit__(self, exc, value, tb):
51        self.cleanup()
52
53    def __del__(self):
54        # Issue a ResourceWarning if implicit cleanup needed
55        self.cleanup()
56
57    # XXX (ncoghlan): The following code attempts to make
58    # this class tolerant of the module nulling out process
59    # that happens during CPython interpreter shutdown
60    # Alas, it doesn't actually manage it. See issue #10188
61    _listdir = staticmethod(_os.listdir)
62    _path_join = staticmethod(_os.path.join)
63    _isdir = staticmethod(_os.path.isdir)
64    _islink = staticmethod(_os.path.islink)
65    _remove = staticmethod(_os.remove)
66    _rmdir = staticmethod(_os.rmdir)
67    _warn = _warnings.warn
68
69    def _rmtree(self, path):
70        # Essentially a stripped down version of shutil.rmtree.  We can't
71        # use globals because they may be None'ed out at shutdown.
72        for name in self._listdir(path):
73            fullname = self._path_join(path, name)
74            try:
75                isdir = self._isdir(fullname) and not self._islink(fullname)
76            except OSError:
77                isdir = False
78            if isdir:
79                self._rmtree(fullname)
80            else:
81                try:
82                    self._remove(fullname)
83                except OSError:
84                    pass
85        try:
86            self._rmdir(path)
87        except OSError:
88            pass
89