123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687 |
- ''' Contexts for *with* statement providing temporary directories
- '''
- from __future__ import division, print_function, absolute_import
- 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)
|