Show More
@@ -1,144 +1,54 b'' | |||
|
1 | """TemporaryDirectory class, copied from Python 3.2. | |
|
2 | ||
|
3 | This is copied from the stdlib and will be standard in Python 3.2 and onwards. | |
|
4 | """ | |
|
5 | ||
|
6 | 1 |
|
|
7 | 2 | import warnings as _warnings |
|
8 | 3 | import sys as _sys |
|
9 | 4 | |
|
10 | # This code should only be used in Python versions < 3.2, since after that we | |
|
11 | # can rely on the stdlib itself. | |
|
12 | try: | |
|
13 | 5 |
|
|
14 | 6 | |
|
15 | except ImportError: | |
|
16 | from tempfile import mkdtemp, template | |
|
17 | ||
|
18 | class TemporaryDirectory(object): | |
|
19 | """Create and return a temporary directory. This has the same | |
|
20 | behavior as mkdtemp but can be used as a context manager. For | |
|
21 | example: | |
|
22 | ||
|
23 | with TemporaryDirectory() as tmpdir: | |
|
24 | ... | |
|
25 | ||
|
26 | Upon exiting the context, the directory and everthing contained | |
|
27 | in it are removed. | |
|
28 | """ | |
|
29 | ||
|
30 | def __init__(self, suffix="", prefix=template, dir=None): | |
|
31 | self.name = mkdtemp(suffix, prefix, dir) | |
|
32 | self._closed = False | |
|
33 | ||
|
34 | def __enter__(self): | |
|
35 | return self.name | |
|
36 | ||
|
37 | def cleanup(self, _warn=False): | |
|
38 | if self.name and not self._closed: | |
|
39 | try: | |
|
40 | self._rmtree(self.name) | |
|
41 | except (TypeError, AttributeError) as ex: | |
|
42 | # Issue #10188: Emit a warning on stderr | |
|
43 | # if the directory could not be cleaned | |
|
44 | # up due to missing globals | |
|
45 | if "None" not in str(ex): | |
|
46 | raise | |
|
47 | print("ERROR: {!r} while cleaning up {!r}".format(ex, self,), | |
|
48 | file=_sys.stderr) | |
|
49 | return | |
|
50 | self._closed = True | |
|
51 | if _warn: | |
|
52 | self._warn("Implicitly cleaning up {!r}".format(self), | |
|
53 | Warning) | |
|
54 | ||
|
55 | def __exit__(self, exc, value, tb): | |
|
56 | self.cleanup() | |
|
57 | ||
|
58 | def __del__(self): | |
|
59 | # Issue a ResourceWarning if implicit cleanup needed | |
|
60 | self.cleanup(_warn=True) | |
|
61 | ||
|
62 | ||
|
63 | # XXX (ncoghlan): The following code attempts to make | |
|
64 | # this class tolerant of the module nulling out process | |
|
65 | # that happens during CPython interpreter shutdown | |
|
66 | # Alas, it doesn't actually manage it. See issue #10188 | |
|
67 | _listdir = staticmethod(_os.listdir) | |
|
68 | _path_join = staticmethod(_os.path.join) | |
|
69 | _isdir = staticmethod(_os.path.isdir) | |
|
70 | _remove = staticmethod(_os.remove) | |
|
71 | _rmdir = staticmethod(_os.rmdir) | |
|
72 | _os_error = _os.error | |
|
73 | _warn = _warnings.warn | |
|
74 | ||
|
75 | def _rmtree(self, path): | |
|
76 | # Essentially a stripped down version of shutil.rmtree. We can't | |
|
77 | # use globals because they may be None'ed out at shutdown. | |
|
78 | for name in self._listdir(path): | |
|
79 | fullname = self._path_join(path, name) | |
|
80 | try: | |
|
81 | isdir = self._isdir(fullname) | |
|
82 | except self._os_error: | |
|
83 | isdir = False | |
|
84 | if isdir: | |
|
85 | self._rmtree(fullname) | |
|
86 | else: | |
|
87 | try: | |
|
88 | self._remove(fullname) | |
|
89 | except self._os_error: | |
|
90 | pass | |
|
91 | try: | |
|
92 | self._rmdir(path) | |
|
93 | except self._os_error: | |
|
94 | pass | |
|
95 | ||
|
96 | ||
|
97 | 7 | class NamedFileInTemporaryDirectory(object): |
|
98 | 8 | |
|
99 | 9 | def __init__(self, filename, mode='w+b', bufsize=-1, **kwds): |
|
100 | 10 | """ |
|
101 | 11 | Open a file named `filename` in a temporary directory. |
|
102 | 12 | |
|
103 | 13 | This context manager is preferred over `NamedTemporaryFile` in |
|
104 | 14 | stdlib `tempfile` when one needs to reopen the file. |
|
105 | 15 | |
|
106 | 16 | Arguments `mode` and `bufsize` are passed to `open`. |
|
107 | 17 | Rest of the arguments are passed to `TemporaryDirectory`. |
|
108 | 18 | |
|
109 | 19 | """ |
|
110 | 20 | self._tmpdir = TemporaryDirectory(**kwds) |
|
111 | 21 | path = _os.path.join(self._tmpdir.name, filename) |
|
112 | 22 | self.file = open(path, mode, bufsize) |
|
113 | 23 | |
|
114 | 24 | def cleanup(self): |
|
115 | 25 | self.file.close() |
|
116 | 26 | self._tmpdir.cleanup() |
|
117 | 27 | |
|
118 | 28 | __del__ = cleanup |
|
119 | 29 | |
|
120 | 30 | def __enter__(self): |
|
121 | 31 | return self.file |
|
122 | 32 | |
|
123 | 33 | def __exit__(self, type, value, traceback): |
|
124 | 34 | self.cleanup() |
|
125 | 35 | |
|
126 | 36 | |
|
127 | 37 | class TemporaryWorkingDirectory(TemporaryDirectory): |
|
128 | 38 | """ |
|
129 | 39 | Creates a temporary directory and sets the cwd to that directory. |
|
130 | 40 | Automatically reverts to previous cwd upon cleanup. |
|
131 | 41 | Usage example: |
|
132 | 42 | |
|
133 | 43 | with TemporaryWorkingDirectory() as tmpdir: |
|
134 | 44 | ... |
|
135 | 45 | """ |
|
136 | 46 | def __enter__(self): |
|
137 | 47 | self.old_wd = _os.getcwd() |
|
138 | 48 | _os.chdir(self.name) |
|
139 | 49 | return super(TemporaryWorkingDirectory, self).__enter__() |
|
140 | 50 | |
|
141 | 51 | def __exit__(self, exc, value, tb): |
|
142 | 52 | _os.chdir(self.old_wd) |
|
143 | 53 | return super(TemporaryWorkingDirectory, self).__exit__(exc, value, tb) |
|
144 | 54 |
General Comments 0
You need to be logged in to leave comments.
Login now