1 files changed, 78 insertions, 0 deletions
diff --git a/tests/compat.py b/tests/compat.py
new file mode 100644
index 0000000..383963a
--- /dev/null
+++ b/tests/compat.py
@@ -0,0 +1,78 @@
+# -*- 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