Fix virtualenv environment detection
There seems to be some combination of factors that cause the string
matching of the VIRTUAL_ENV envvar to fail on trybots/pre-cq. Given
that, I think its generally a bad idea to do path comparisons against
the virtualenv environment and reverting back to a dumber, but more
robust implementation of the virtualenv check using a magic variable.
BUG=chromium:673926
TEST=Run bin/venv_check and unit tests
Change-Id: Ifaafafc0abf5b079585a168f5d6a4b731c1a2d2e
Reviewed-on: https://chromium-review.googlesource.com/419722
Commit-Ready: David Riley <davidriley@chromium.org>
Tested-by: Allen Li <ayatane@chromium.org>
Reviewed-by: Aviv Keshet <akeshet@chromium.org>
diff --git a/scripts/virtualenv_wrapper.py b/scripts/virtualenv_wrapper.py
index 61bb7b0..15def65 100755
--- a/scripts/virtualenv_wrapper.py
+++ b/scripts/virtualenv_wrapper.py
@@ -13,6 +13,7 @@
import wrapper
+
def _FindChromiteDir():
path = os.path.dirname(os.path.realpath(__file__))
while not os.path.exists(os.path.join(path, 'PRESUBMIT.cfg')):
@@ -21,15 +22,21 @@
_CHROMITE_DIR = _FindChromiteDir()
+
# _VIRTUALENV_DIR contains the scripts for working with venvs
_VIRTUALENV_DIR = os.path.join(_CHROMITE_DIR, '../infra_virtualenv')
-# _VENV_DIR is the actual virtualenv that contains bin/activate.
+_CREATE_VENV_PATH = os.path.join(_VIRTUALENV_DIR, 'create_venv')
+_VENV_COMMAND_PATH = os.path.join(_VIRTUALENV_DIR, 'venv_command')
+
+# _VENV_DIR is the virtualenv dir that contains bin/activate.
_VENV_DIR = os.path.join(_CHROMITE_DIR, '.venv')
_REQUIREMENTS = os.path.join(_CHROMITE_DIR, 'venv', 'requirements.txt')
+_VENV_MARKER = 'INSIDE_CHROMITE_VENV'
+
def main():
- if _IsInsideVenv():
+ if _IsInsideVenv(os.environ):
wrapper.DoMain()
else:
_CreateVenv()
@@ -39,7 +46,7 @@
def _CreateVenv():
"""Create or update chromite venv."""
subprocess.check_call([
- os.path.join(_VIRTUALENV_DIR, 'create_venv'),
+ _CREATE_VENV_PATH,
_VENV_DIR,
_REQUIREMENTS,
], stdout=sys.stderr)
@@ -51,15 +58,43 @@
Args:
args: Sequence of arguments.
"""
- os.execv(os.path.join(_VIRTUALENV_DIR, 'venv_command'),
- ['venv_command', _VENV_DIR] + list(args))
+ os.execve(_VENV_COMMAND_PATH,
+ [os.path.basename(_VENV_COMMAND_PATH), _VENV_DIR] + list(args),
+ _CreateVenvEnvironment(os.environ))
-def _IsInsideVenv():
- """Return whether we're running inside a virtualenv."""
- # Proper way is checking sys.prefix and sys.base_prefix in Python 3.
- # PEP 405 isn't fully implemented in Python 2.
- return _VENV_DIR in os.environ.get('VIRTUAL_ENV', '')
+def _CreateVenvEnvironment(env_dict):
+ """Create environment for a virtualenv.
+
+ This adds a special marker variable to a copy of the input environment dict
+ and returns the copy.
+
+ Args:
+ env_dict: Environment variable dict to use as base, which is not modified.
+
+ Returns:
+ New environment dict for a virtualenv.
+ """
+ new_env_dict = env_dict.copy()
+ new_env_dict[_VENV_MARKER] = '1'
+ return new_env_dict
+
+
+def _IsInsideVenv(env_dict):
+ """Return whether the environment dict is running inside a virtualenv.
+
+ This checks the environment dict for the special marker added by
+ _CreateVenvEnvironment().
+
+ Args:
+ env_dict: Environment variable dict to check
+
+ Returns:
+ A true value if inside virtualenv, else a false value.
+ """
+ # Checking sys.prefix or doing any kind of path check is unreliable because
+ # we check out chromite to weird places.
+ return env_dict.get(_VENV_MARKER, '')
if __name__ == '__main__':