Coverage for gwcelery/util/tempfile.py: 100%

1import tempfile 

2from contextlib import contextmanager 

3 

4__all__ = ('NamedTemporaryFile',) 

5 

6 

7@contextmanager 

8def NamedTemporaryFile(content=None, **kwargs): # noqa: N802 

9 r"""Convenience wrapper for :func:`tempfile.NamedTemporaryFile` that writes 

10 some data to the file before handing it to the calling code. 

11 

12 Parameters 

13 ---------- 

14 content : str, bytes, None 

15 Initial contents of the file. 

16 \**kwargs 

17 Additional keyword arguments to pass to 

18 :func:`tempfile.NamedTemporaryFile`. 

19 

20 """ 

21 if isinstance(content, bytes): 

22 kwargs = dict(kwargs, mode='w+b') 

23 elif isinstance(content, str): 

24 kwargs = dict(kwargs, mode='w+') 

25 elif content is not None: 

26 raise TypeError('content is of unknown type') 

27 with tempfile.NamedTemporaryFile(**kwargs) as f: 

28 if content is not None: 

29 f.write(content) 

30 f.flush() 

31 f.seek(0) 

32 yield f