|
|
''' 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) |
|
|
|
