| ''' Contexts for *with* statement providing temporary directories |
| ''' |
| import os |
| from contextlib import contextmanager |
| from shutil import rmtree |
| from tempfile import mkdtemp |
|
|
|
|
| @contextmanager |
| def tempdir(): |
| """Create and return a temporary directory. This has the same |
| behavior as mkdtemp but can be used as a context manager. |
| |
| Upon exiting the context, the directory and everything contained |
| in it are removed. |
| |
| Examples |
| -------- |
| >>> import os |
| >>> with tempdir() as tmpdir: |
| ... fname = os.path.join(tmpdir, 'example_file.txt') |
| ... with open(fname, 'wt') as fobj: |
| ... _ = fobj.write('a string\\n') |
| >>> os.path.exists(tmpdir) |
| False |
| """ |
| d = mkdtemp() |
| yield d |
| rmtree(d) |
|
|
|
|
| @contextmanager |
| def in_tempdir(): |
| ''' Create, return, and change directory to a temporary directory |
| |
| Examples |
| -------- |
| >>> import os |
| >>> my_cwd = os.getcwd() |
| >>> with in_tempdir() as tmpdir: |
| ... _ = open('test.txt', 'wt').write('some text') |
| ... assert os.path.isfile('test.txt') |
| ... assert os.path.isfile(os.path.join(tmpdir, 'test.txt')) |
| >>> os.path.exists(tmpdir) |
| False |
| >>> os.getcwd() == my_cwd |
| True |
| ''' |
| pwd = os.getcwd() |
| d = mkdtemp() |
| os.chdir(d) |
| yield d |
| os.chdir(pwd) |
| rmtree(d) |
|
|
|
|
| @contextmanager |
| def in_dir(dir=None): |
| """ Change directory to given directory for duration of ``with`` block |
| |
| Useful when you want to use `in_tempdir` for the final test, but |
| you are still debugging. For example, you may want to do this in the end: |
| |
| >>> with in_tempdir() as tmpdir: |
| ... # do something complicated which might break |
| ... pass |
| |
| But, indeed, the complicated thing does break, and meanwhile, the |
| ``in_tempdir`` context manager wiped out the directory with the |
| temporary files that you wanted for debugging. So, while debugging, you |
| replace with something like: |
| |
| >>> with in_dir() as tmpdir: # Use working directory by default |
| ... # do something complicated which might break |
| ... pass |
| |
| You can then look at the temporary file outputs to debug what is happening, |
| fix, and finally replace ``in_dir`` with ``in_tempdir`` again. |
| """ |
| cwd = os.getcwd() |
| if dir is None: |
| yield cwd |
| return |
| os.chdir(dir) |
| yield dir |
| os.chdir(cwd) |
|
|
