API: Allow passing on the environment variable as an argument (#2054)

Defaults to os.environ.

Signed-off-by: Bernát Gábor <bgabor8@bloomberg.net>

docs/changelog/2054.feature.rst

@@ -0,0 +1,2 @@
+On the programmatic API allow passing in the environment variable dictionary to use, defaults to ``os.environ`` if not
+specified - by :user:`gaborbernat`.

src/virtualenv/__main__.py

@@ -1,19 +1,21 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
 import logging
+import os
 import sys
 from datetime import datetime
 
 
-def run(args=None, options=None):
+def run(args=None, options=None, env=None):
+    env = os.environ if env is None else env
     start = datetime.now()
     from virtualenv.run import cli_run
     from virtualenv.util.error import ProcessCallFailed
 
     if args is None:
         args = sys.argv[1:]
     try:
-        session = cli_run(args, options)
+        session = cli_run(args, options, env)
         logging.warning(LogSession(session, start))
     except ProcessCallFailed as exception:
         print("subprocess call failed for {} with code {}".format(exception.cmd, exception.code))
@@ -54,12 +56,13 @@ def __str__(self):
         return "\n".join(lines)
 
 
-def run_with_catch(args=None):
+def run_with_catch(args=None, env=None):
     from virtualenv.config.cli.parser import VirtualEnvOptions
 
+    env = os.environ if env is None else env
     options = VirtualEnvOptions()
     try:
-        run(args, options)
+        run(args, options, env)
     except (KeyboardInterrupt, SystemExit, Exception) as exception:
         try:
             if getattr(options, "with_traceback", False):

src/virtualenv/app_data/__init__.py

@@ -14,21 +14,22 @@
 from .via_tempdir import TempAppData
 
 
-def _default_app_data_dir():  # type: () -> str
+def _default_app_data_dir(env):
     key = str("VIRTUALENV_OVERRIDE_APP_DATA")
-    if key in os.environ:
-        return os.environ[key]
+    if key in env:
+        return env[key]
     else:
         return user_data_dir(appname="virtualenv", appauthor="pypa")
 
 
 def make_app_data(folder, **kwargs):
     read_only = kwargs.pop("read_only")
+    env = kwargs.pop("env")
     if kwargs:  # py3+ kwonly
         raise TypeError("unexpected keywords: {}")
 
     if folder is None:
-        folder = _default_app_data_dir()
+        folder = _default_app_data_dir(env)
     folder = os.path.abspath(folder)
 
     if read_only:

src/virtualenv/config/cli/parser.py

@@ -1,5 +1,6 @@
 from __future__ import absolute_import, unicode_literals
 
+import os
 from argparse import SUPPRESS, ArgumentDefaultsHelpFormatter, ArgumentParser, Namespace
 from collections import OrderedDict
 
@@ -47,9 +48,11 @@ class VirtualEnvConfigParser(ArgumentParser):
     Custom option parser which updates its defaults by checking the configuration files and environmental variables
     """
 
-    def __init__(self, options=None, *args, **kwargs):
-        self.file_config = IniConfig()
+    def __init__(self, options=None, env=None, *args, **kwargs):
+        env = os.environ if env is None else env
+        self.file_config = IniConfig(env)
         self.epilog_list = []
+        self.env = env
         kwargs["epilog"] = self.file_config.epilog
         kwargs["add_help"] = False
         kwargs["formatter_class"] = HelpFormatter
@@ -75,7 +78,7 @@ def _fix_default(self, action):
             names = OrderedDict((i.lstrip("-").replace("-", "_"), None) for i in action.option_strings)
             outcome = None
             for name in names:
-                outcome = get_env_var(name, as_type)
+                outcome = get_env_var(name, as_type, self.env)
                 if outcome is not None:
                     break
             if outcome is None and self.file_config:
@@ -101,6 +104,7 @@ def parse_known_args(self, args=None, namespace=None):
         self._fix_defaults()
         self.options._src = "cli"
         try:
+            namespace.env = self.env
             return super(VirtualEnvConfigParser, self).parse_known_args(args, namespace=namespace)
         finally:
             self.options._src = None

src/virtualenv/config/env_var.py

@@ -1,22 +1,21 @@
 from __future__ import absolute_import, unicode_literals
 
-import os
-
 from virtualenv.util.six import ensure_str, ensure_text
 
 from .convert import convert
 
 
-def get_env_var(key, as_type):
+def get_env_var(key, as_type, env):
     """Get the environment variable option.
 
     :param key: the config key requested
     :param as_type: the type we would like to convert it to
+    :param env: environment variables to use
     :return:
     """
     environ_key = ensure_str("VIRTUALENV_{}".format(key.upper()))
-    if os.environ.get(environ_key):
-        value = os.environ[environ_key]
+    if env.get(environ_key):
+        value = env[environ_key]
         # noinspection PyBroadException
         try:
             source = "env var {}".format(ensure_text(environ_key))

src/virtualenv/config/ini.py

@@ -19,8 +19,9 @@ class IniConfig(object):
 
     section = "virtualenv"
 
-    def __init__(self):
-        config_file = os.environ.get(self.VIRTUALENV_CONFIG_FILE_ENV_VAR, None)
+    def __init__(self, env=None):
+        env = os.environ if env is None else env
+        config_file = env.get(self.VIRTUALENV_CONFIG_FILE_ENV_VAR, None)
         self.is_env_var = config_file is not None
         config_file = (
             Path(config_file)

src/virtualenv/create/creator.py

@@ -47,6 +47,7 @@ def __init__(self, options, interpreter):
         self.no_vcs_ignore = options.no_vcs_ignore
         self.pyenv_cfg = PyEnvCfg.from_folder(self.dest)
         self.app_data = options.app_data
+        self.env = options.env
 
     def __repr__(self):
         return ensure_str(self.__unicode__())
@@ -204,16 +205,16 @@ def debug(self):
         :return: debug information about the virtual environment (only valid after :meth:`create` has run)
         """
         if self._debug is None and self.exe is not None:
-            self._debug = get_env_debug_info(self.exe, self.debug_script(), self.app_data)
+            self._debug = get_env_debug_info(self.exe, self.debug_script(), self.app_data, self.env)
         return self._debug
 
     # noinspection PyMethodMayBeStatic
     def debug_script(self):
         return DEBUG_SCRIPT
 
 
-def get_env_debug_info(env_exe, debug_script, app_data):
-    env = os.environ.copy()
+def get_env_debug_info(env_exe, debug_script, app_data, env):
+    env = env.copy()
     env.pop(str("PYTHONPATH"), None)
 
     with app_data.ensure_extracted(debug_script) as debug_script:

src/virtualenv/discovery/builtin.py

@@ -44,7 +44,7 @@ def add_parser_arguments(cls, parser):
 
     def run(self):
         for python_spec in self.python_spec:
-            result = get_interpreter(python_spec, self.try_first_with, self.app_data)
+            result = get_interpreter(python_spec, self.try_first_with, self.app_data, self._env)
             if result is not None:
                 return result
         return None
@@ -57,11 +57,12 @@ def __unicode__(self):
         return "{} discover of python_spec={!r}".format(self.__class__.__name__, spec)
 
 
-def get_interpreter(key, try_first_with, app_data=None):
+def get_interpreter(key, try_first_with, app_data=None, env=None):
     spec = PythonSpec.from_string_spec(key)
     logging.info("find interpreter for spec %r", spec)
     proposed_paths = set()
-    for interpreter, impl_must_match in propose_interpreters(spec, try_first_with, app_data):
+    env = os.environ if env is None else env
+    for interpreter, impl_must_match in propose_interpreters(spec, try_first_with, app_data, env):
         key = interpreter.system_executable, impl_must_match
         if key in proposed_paths:
             continue
@@ -72,16 +73,17 @@ def get_interpreter(key, try_first_with, app_data=None):
         proposed_paths.add(key)
 
 
-def propose_interpreters(spec, try_first_with, app_data):
+def propose_interpreters(spec, try_first_with, app_data, env=None):
     # 0. try with first
+    env = os.environ if env is None else env
     for py_exe in try_first_with:
         path = os.path.abspath(py_exe)
         try:
             os.lstat(path)  # Windows Store Python does not work with os.path.exists, but does for os.lstat
         except OSError:
             pass
         else:
-            yield PythonInfo.from_exe(os.path.abspath(path), app_data), True
+            yield PythonInfo.from_exe(os.path.abspath(path), app_data, env=env), True
 
     # 1. if it's a path and exists
     if spec.path is not None:
@@ -91,7 +93,7 @@ def propose_interpreters(spec, try_first_with, app_data):
             if spec.is_abs:
                 raise
         else:
-            yield PythonInfo.from_exe(os.path.abspath(spec.path), app_data), True
+            yield PythonInfo.from_exe(os.path.abspath(spec.path), app_data, env=env), True
         if spec.is_abs:
             return
     else:
@@ -102,27 +104,27 @@ def propose_interpreters(spec, try_first_with, app_data):
         if IS_WIN:
             from .windows import propose_interpreters
 
-            for interpreter in propose_interpreters(spec, app_data):
+            for interpreter in propose_interpreters(spec, app_data, env):
                 yield interpreter, True
     # finally just find on path, the path order matters (as the candidates are less easy to control by end user)
-    paths = get_paths()
+    paths = get_paths(env)
     tested_exes = set()
     for pos, path in enumerate(paths):
         path = ensure_text(path)
-        logging.debug(LazyPathDump(pos, path))
+        logging.debug(LazyPathDump(pos, path, env))
         for candidate, match in possible_specs(spec):
             found = check_path(candidate, path)
             if found is not None:
                 exe = os.path.abspath(found)
                 if exe not in tested_exes:
                     tested_exes.add(exe)
-                    interpreter = PathPythonInfo.from_exe(exe, app_data, raise_on_error=False)
+                    interpreter = PathPythonInfo.from_exe(exe, app_data, raise_on_error=False, env=env)
                     if interpreter is not None:
                         yield interpreter, match
 
 
-def get_paths():
-    path = os.environ.get(str("PATH"), None)
+def get_paths(env):
+    path = env.get(str("PATH"), None)
     if path is None:
         try:
             path = os.confstr("CS_PATH")
@@ -136,16 +138,17 @@ def get_paths():
 
 
 class LazyPathDump(object):
-    def __init__(self, pos, path):
+    def __init__(self, pos, path, env):
         self.pos = pos
         self.path = path
+        self.env = env
 
     def __repr__(self):
         return ensure_str(self.__unicode__())
 
     def __unicode__(self):
         content = "discover PATH[{}]={}".format(self.pos, self.path)
-        if os.environ.get(str("_VIRTUALENV_DEBUG")):  # this is the over the board debug
+        if self.env.get(str("_VIRTUALENV_DEBUG")):  # this is the over the board debug
             content += " with =>"
             for file_name in os.listdir(self.path):
                 try:

src/virtualenv/discovery/cached_py_info.py

@@ -23,9 +23,9 @@
 _CACHE[Path(sys.executable)] = PythonInfo()
 
 
-def from_exe(cls, app_data, exe, raise_on_error=True, ignore_cache=False):
-    """"""
-    result = _get_from_cache(cls, app_data, exe, ignore_cache=ignore_cache)
+def from_exe(cls, app_data, exe, env=None, raise_on_error=True, ignore_cache=False):
+    env = os.environ if env is None else env
+    result = _get_from_cache(cls, app_data, exe, env, ignore_cache=ignore_cache)
     if isinstance(result, Exception):
         if raise_on_error:
             raise result
@@ -35,22 +35,22 @@ def from_exe(cls, app_data, exe, raise_on_error=True, ignore_cache=False):
     return result
 
 
-def _get_from_cache(cls, app_data, exe, ignore_cache=True):
+def _get_from_cache(cls, app_data, exe, env, ignore_cache=True):
     # note here we cannot resolve symlinks, as the symlink may trigger different prefix information if there's a
     # pyenv.cfg somewhere alongside on python3.4+
     exe_path = Path(exe)
     if not ignore_cache and exe_path in _CACHE:  # check in the in-memory cache
         result = _CACHE[exe_path]
     else:  # otherwise go through the app data cache
-        py_info = _get_via_file_cache(cls, app_data, exe_path, exe)
+        py_info = _get_via_file_cache(cls, app_data, exe_path, exe, env)
         result = _CACHE[exe_path] = py_info
     # independent if it was from the file or in-memory cache fix the original executable location
     if isinstance(result, PythonInfo):
         result.executable = exe
     return result
 
 
-def _get_via_file_cache(cls, app_data, path, exe):
+def _get_via_file_cache(cls, app_data, path, exe, env):
     path_text = ensure_text(str(path))
     try:
         path_modified = path.stat().st_mtime
@@ -72,7 +72,7 @@ def _get_via_file_cache(cls, app_data, path, exe):
             else:
                 py_info_store.remove()
         if py_info is None:  # if not loaded run and save
-            failure, py_info = _run_subprocess(cls, exe, app_data)
+            failure, py_info = _run_subprocess(cls, exe, app_data, env)
             if failure is None:
                 data = {"st_mtime": path_modified, "path": path_text, "content": py_info._to_dict()}
                 py_info_store.write(data)
@@ -81,12 +81,12 @@ def _get_via_file_cache(cls, app_data, path, exe):
     return py_info
 
 
-def _run_subprocess(cls, exe, app_data):
+def _run_subprocess(cls, exe, app_data, env):
     py_info_script = Path(os.path.abspath(__file__)).parent / "py_info.py"
     with app_data.ensure_extracted(py_info_script) as py_info_script:
         cmd = [exe, str(py_info_script)]
         # prevent sys.prefix from leaking into the child process - see https://bugs.python.org/issue22490
-        env = os.environ.copy()
+        env = env.copy()
         env.pop("__PYVENV_LAUNCHER__", None)
         logging.debug("get interpreter info via cmd: %s", LogCmd(cmd))
         try:

src/virtualenv/discovery/discover.py

@@ -25,6 +25,7 @@ def __init__(self, options):
         """
         self._has_run = False
         self._interpreter = None
+        self._env = options.env
 
     @abstractmethod
     def run(self):