1# -*- coding: utf-8 -*- 2"""Python compatibility support for CLI Helpers' tests.""" 3 4from __future__ import unicode_literals 5import os as _os 6import shutil as _shutil 7import tempfile as _tempfile 8import warnings as _warnings 9 10from cli_helpers.compat import PY2 11 12 13class _TempDirectory(object): 14 """Create and return a temporary directory. This has the same 15 behavior as mkdtemp but can be used as a context manager. For 16 example: 17 18 with TemporaryDirectory() as tmpdir: 19 ... 20 21 Upon exiting the context, the directory and everything contained 22 in it are removed. 23 24 NOTE: Copied from the Python 3 standard library. 25 """ 26 27 # Handle mkdtemp raising an exception 28 name = None 29 _closed = False 30 31 def __init__(self, suffix="", prefix="tmp", dir=None): 32 self.name = _tempfile.mkdtemp(suffix, prefix, dir) 33 34 def __repr__(self): 35 return "<{} {!r}>".format(self.__class__.__name__, self.name) 36 37 def __enter__(self): 38 return self.name 39 40 def cleanup(self, _warn=False, _warnings=_warnings): 41 if self.name and not self._closed: 42 try: 43 _shutil.rmtree(self.name) 44 except (TypeError, AttributeError) as ex: 45 if "None" not in "%s" % (ex,): 46 raise 47 self._rmtree(self.name) 48 self._closed = True 49 if _warn and _warnings.warn: 50 _warnings.warn( 51 "Implicitly cleaning up {!r}".format(self), ResourceWarning 52 ) 53 54 def __exit__(self, exc, value, tb): 55 self.cleanup() 56 57 def __del__(self): 58 # Issue a ResourceWarning if implicit cleanup needed 59 self.cleanup(_warn=True) 60 61 def _rmtree( 62 self, 63 path, 64 _OSError=OSError, 65 _sep=_os.path.sep, 66 _listdir=_os.listdir, 67 _remove=_os.remove, 68 _rmdir=_os.rmdir, 69 ): 70 # Essentially a stripped down version of shutil.rmtree. We can't 71 # use globals because they may be None'ed out at shutdown. 72 if not isinstance(path, str): 73 _sep = _sep.encode() 74 try: 75 for name in _listdir(path): 76 fullname = path + _sep + name 77 try: 78 _remove(fullname) 79 except _OSError: 80 self._rmtree(fullname) 81 _rmdir(path) 82 except _OSError: 83 pass 84 85 86TemporaryDirectory = _TempDirectory if PY2 else _tempfile.TemporaryDirectory 87