|
| 1 | +import contextlib |
| 2 | +import logging |
| 3 | +import os |
| 4 | +import subprocess |
| 5 | +import shlex |
| 6 | +import sys |
| 7 | +import sysconfig |
| 8 | +import tempfile |
| 9 | +import venv |
| 10 | + |
| 11 | + |
| 12 | +class VirtualEnvironment: |
| 13 | + def __init__(self, prefix, **venv_create_args): |
| 14 | + self._logger = logging.getLogger(self.__class__.__name__) |
| 15 | + venv.create(prefix, **venv_create_args) |
| 16 | + self._prefix = prefix |
| 17 | + self._paths = sysconfig.get_paths( |
| 18 | + scheme='venv', |
| 19 | + vars={'base': self.prefix}, |
| 20 | + expand=True, |
| 21 | + ) |
| 22 | + |
| 23 | + @classmethod |
| 24 | + @contextlib.contextmanager |
| 25 | + def from_tmpdir(cls, *, prefix=None, dir=None, **venv_create_args): |
| 26 | + delete = not bool(os.environ.get('PYTHON_TESTS_KEEP_VENV')) |
| 27 | + with tempfile.TemporaryDirectory(prefix=prefix, dir=dir, delete=delete) as tmpdir: |
| 28 | + yield cls(tmpdir, **venv_create_args) |
| 29 | + |
| 30 | + @property |
| 31 | + def prefix(self): |
| 32 | + return self._prefix |
| 33 | + |
| 34 | + @property |
| 35 | + def paths(self): |
| 36 | + return self._paths |
| 37 | + |
| 38 | + @property |
| 39 | + def interpreter(self): |
| 40 | + return os.path.join(self.paths['scripts'], os.path.basename(sys.executable)) |
| 41 | + |
| 42 | + def _format_output(self, name, data, indent='\t'): |
| 43 | + if not data: |
| 44 | + return indent + f'{name}: (none)' |
| 45 | + if len(data.splitlines()) == 1: |
| 46 | + return indent + f'{name}: {data}' |
| 47 | + else: |
| 48 | + prefixed_lines = '\n'.join(indent + '> ' + line for line in data.splitlines()) |
| 49 | + return indent + f'{name}:\n' + prefixed_lines |
| 50 | + |
| 51 | + def run(self, *args, **subprocess_args): |
| 52 | + if subprocess_args.get('shell'): |
| 53 | + raise ValueError('Running the subprocess in shell mode is not supported.') |
| 54 | + default_args = { |
| 55 | + 'capture_output': True, |
| 56 | + 'check': True, |
| 57 | + } |
| 58 | + try: |
| 59 | + result = subprocess.run([self.interpreter, *args], **default_args | subprocess_args) |
| 60 | + except subprocess.CalledProcessError as e: |
| 61 | + if e.returncode != 0: |
| 62 | + self._logger.error( |
| 63 | + f'Interpreter returned non-zero exit status {e.returncode}.\n' |
| 64 | + + self._format_output('COMMAND', shlex.join(e.cmd)) + '\n' |
| 65 | + + self._format_output('STDOUT', e.stdout.decode()) + '\n' |
| 66 | + + self._format_output('STDERR', e.stderr.decode()) + '\n' |
| 67 | + ) |
| 68 | + raise |
| 69 | + else: |
| 70 | + return result |
0 commit comments