1import contextlib 2import logging 3import os 4import subprocess 5import shlex 6import sys 7import sysconfig 8import tempfile 9import venv 10 11 12class 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 71