Show More
@@ -1,54 +1,60 b'' | |||||
|
1 | """ This module contains classes - NamedFileInTemporaryDirectory, TemporaryWorkingDirectory. | |||
|
2 | ||||
|
3 | These classes add extra features such as creating a named file in temporary directory and | |||
|
4 | creating a context manager for the working directory which is also temporary. | |||
|
5 | """ | |||
|
6 | ||||
1 | import os as _os |
|
7 | import os as _os | |
2 | import warnings as _warnings |
|
8 | import warnings as _warnings | |
3 | import sys as _sys |
|
9 | import sys as _sys | |
4 |
|
10 | |||
5 | from tempfile import TemporaryDirectory |
|
11 | from tempfile import TemporaryDirectory | |
6 |
|
12 | |||
7 | class NamedFileInTemporaryDirectory(object): |
|
13 | class NamedFileInTemporaryDirectory(object): | |
8 |
|
14 | |||
9 | def __init__(self, filename, mode='w+b', bufsize=-1, **kwds): |
|
15 | def __init__(self, filename, mode='w+b', bufsize=-1, **kwds): | |
10 | """ |
|
16 | """ | |
11 | Open a file named `filename` in a temporary directory. |
|
17 | Open a file named `filename` in a temporary directory. | |
12 |
|
18 | |||
13 | This context manager is preferred over `NamedTemporaryFile` in |
|
19 | This context manager is preferred over `NamedTemporaryFile` in | |
14 | stdlib `tempfile` when one needs to reopen the file. |
|
20 | stdlib `tempfile` when one needs to reopen the file. | |
15 |
|
21 | |||
16 | Arguments `mode` and `bufsize` are passed to `open`. |
|
22 | Arguments `mode` and `bufsize` are passed to `open`. | |
17 | Rest of the arguments are passed to `TemporaryDirectory`. |
|
23 | Rest of the arguments are passed to `TemporaryDirectory`. | |
18 |
|
24 | |||
19 | """ |
|
25 | """ | |
20 | self._tmpdir = TemporaryDirectory(**kwds) |
|
26 | self._tmpdir = TemporaryDirectory(**kwds) | |
21 | path = _os.path.join(self._tmpdir.name, filename) |
|
27 | path = _os.path.join(self._tmpdir.name, filename) | |
22 | self.file = open(path, mode, bufsize) |
|
28 | self.file = open(path, mode, bufsize) | |
23 |
|
29 | |||
24 | def cleanup(self): |
|
30 | def cleanup(self): | |
25 | self.file.close() |
|
31 | self.file.close() | |
26 | self._tmpdir.cleanup() |
|
32 | self._tmpdir.cleanup() | |
27 |
|
33 | |||
28 | __del__ = cleanup |
|
34 | __del__ = cleanup | |
29 |
|
35 | |||
30 | def __enter__(self): |
|
36 | def __enter__(self): | |
31 | return self.file |
|
37 | return self.file | |
32 |
|
38 | |||
33 | def __exit__(self, type, value, traceback): |
|
39 | def __exit__(self, type, value, traceback): | |
34 | self.cleanup() |
|
40 | self.cleanup() | |
35 |
|
41 | |||
36 |
|
42 | |||
37 | class TemporaryWorkingDirectory(TemporaryDirectory): |
|
43 | class TemporaryWorkingDirectory(TemporaryDirectory): | |
38 | """ |
|
44 | """ | |
39 | Creates a temporary directory and sets the cwd to that directory. |
|
45 | Creates a temporary directory and sets the cwd to that directory. | |
40 | Automatically reverts to previous cwd upon cleanup. |
|
46 | Automatically reverts to previous cwd upon cleanup. | |
41 | Usage example: |
|
47 | Usage example: | |
42 |
|
48 | |||
43 | with TemporaryWorkingDirectory() as tmpdir: |
|
49 | with TemporaryWorkingDirectory() as tmpdir: | |
44 | ... |
|
50 | ... | |
45 | """ |
|
51 | """ | |
46 | def __enter__(self): |
|
52 | def __enter__(self): | |
47 | self.old_wd = _os.getcwd() |
|
53 | self.old_wd = _os.getcwd() | |
48 | _os.chdir(self.name) |
|
54 | _os.chdir(self.name) | |
49 | return super(TemporaryWorkingDirectory, self).__enter__() |
|
55 | return super(TemporaryWorkingDirectory, self).__enter__() | |
50 |
|
56 | |||
51 | def __exit__(self, exc, value, tb): |
|
57 | def __exit__(self, exc, value, tb): | |
52 | _os.chdir(self.old_wd) |
|
58 | _os.chdir(self.old_wd) | |
53 | return super(TemporaryWorkingDirectory, self).__exit__(exc, value, tb) |
|
59 | return super(TemporaryWorkingDirectory, self).__exit__(exc, value, tb) | |
54 |
|
60 |
General Comments 0
You need to be logged in to leave comments.
Login now