# -*- coding: utf-8 -*- """Python compatibility support for CLI Helpers' tests.""" from __future__ import unicode_literals import os as _os import shutil as _shutil import tempfile as _tempfile import warnings as _warnings from cli_helpers.compat import PY2 class _TempDirectory(object): """Create and return a temporary directory. This has the same behavior as mkdtemp but can be used as a context manager. For example: with TemporaryDirectory() as tmpdir: ... Upon exiting the context, the directory and everything contained in it are removed. NOTE: Copied from the Python 3 standard library. """ # Handle mkdtemp raising an exception name = None _closed = False def __init__(self, suffix="", prefix='tmp', dir=None): self.name = _tempfile.mkdtemp(suffix, prefix, dir) def __repr__(self): return "<{} {!r}>".format(self.__class__.__name__, self.name) def __enter__(self): return self.name def cleanup(self, _warn=False, _warnings=_warnings): if self.name and not self._closed: try: _shutil.rmtree(self.name) except (TypeError, AttributeError) as ex: if "None" not in '%s' % (ex,): raise self._rmtree(self.name) self._closed = True if _warn and _warnings.warn: _warnings.warn("Implicitly cleaning up {!r}".format(self), ResourceWarning) def __exit__(self, exc, value, tb): self.cleanup() def __del__(self): # Issue a ResourceWarning if implicit cleanup needed self.cleanup(_warn=True) def _rmtree(self, path, _OSError=OSError, _sep=_os.path.sep, _listdir=_os.listdir, _remove=_os.remove, _rmdir=_os.rmdir): # Essentially a stripped down version of shutil.rmtree. We can't # use globals because they may be None'ed out at shutdown. if not isinstance(path, str): _sep = _sep.encode() try: for name in _listdir(path): fullname = path + _sep + name try: _remove(fullname) except _OSError: self._rmtree(fullname) _rmdir(path) except _OSError: pass TemporaryDirectory = _TempDirectory if PY2 else _tempfile.TemporaryDirectory