Skip to content

Instantly share code, notes, and snippets.

@waylan
Created May 24, 2018 20:27
Show Gist options
  • Save waylan/50a067d79386aabf9073dbab7beb2950 to your computer and use it in GitHub Desktop.
Save waylan/50a067d79386aabf9073dbab7beb2950 to your computer and use it in GitHub Desktop.
A decorator for building a temporary directory with prepopulated files.
from functools import wraps
try:
# py>=3.2
from tempfile import TemporaryDirectory
except ImportError:
from backports.tempfile import TemporaryDirectory
from pathlib2 import Path
def tempdir(files=None, **kw):
"""
A decorator for building a temporary directory with prepopulated files.
The temproary directory and files are created just before the wrapped function is called and are destroyed
imediately after the wrapped function returns.
The `files` keyword should be a dict of file paths as keys and strings of file content as values.
If `files` is a list, then each item is assumed to be a path of an empty file. All other
keywords are passed to `tempfile.TemporaryDirectory` to create the parent directory.
In the following example, two files are created in the temporary directory and then are destroyed when
the function exits:
@tempdir(files={
'foo.txt': 'foo content',
'bar.txt': 'bar content'
})
def example(dir):
assert Path(dir, 'foo.txt').is_file()
p = Path(dir, 'bar.txt')
assert p.is_file()
assert p.read_text(encoding='utf-8') == 'bar content'
"""
files = {f: '' for f in files} if isinstance(files, (list, tuple)) else files or {}
def decorator(fn):
@wraps(fn)
def wrapper(*args, **kwargs):
with TemporaryDirectory(**kw) as td:
for path, content in files.items():
Path(td, path).write_text(content, encoding='utf-8')
return fn(td, *args, **kwargs)
return wrapper
return decorator
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment