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