Migrate virtualenv wrapper to virtualenv repo
BUG=chromium:645611
TEST=Run venv_check and virtualenv_wrapper_unittest
Change-Id: Ie5471bcb3f3fd02d595f36b7082859c9eed889cc
Reviewed-on: https://chromium-review.googlesource.com/415467
Commit-Ready: Allen Li <ayatane@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 cf74a6e..61bb7b0 100755
--- a/scripts/virtualenv_wrapper.py
+++ b/scripts/virtualenv_wrapper.py
@@ -8,37 +8,59 @@
from __future__ import print_function
import os
+import subprocess
import sys
+
import wrapper
-# TODO(akeshet): This transitively imports a bunch of other dependencies,
-# including cbuildbot.contants. Ideally we wouldn't need that much junk in this
-# wrapper, and importing all that prior to entering the virtualenv might
-# actually cause issues.
-from chromite.lib import cros_build_lib
+def _FindChromiteDir():
+ path = os.path.dirname(os.path.realpath(__file__))
+ while not os.path.exists(os.path.join(path, 'PRESUBMIT.cfg')):
+ path = os.path.dirname(path)
+ return path
-# TODO(akeshet): Since we are using the above lib which imports
-# cbuildbot.constants anyway, we might as well make use of it in determining
-# CHROMITE_PATH. If we want to eliminate this import, we can duplicate the
-# chromite path finding code. It would look something like this:
-# path = os.path.dirname(os.path.realpath(__file__))
-# while not os.path.exists(os.path.join(path, 'PRESUBMIT.cfg')):
-# path = os.path.dirname(path)
-from chromite.lib import constants
-_CHROMITE_DIR = constants.CHROMITE_DIR
-_IN_VENV = 'IN_CHROMITE_VENV'
+_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.
+_VENV_DIR = os.path.join(_CHROMITE_DIR, '.venv')
+_REQUIREMENTS = os.path.join(_CHROMITE_DIR, 'venv', 'requirements.txt')
+
+
+def main():
+ if _IsInsideVenv():
+ wrapper.DoMain()
+ else:
+ _CreateVenv()
+ _ExecInVenv(sys.argv)
+
+
+def _CreateVenv():
+ """Create or update chromite venv."""
+ subprocess.check_call([
+ os.path.join(_VIRTUALENV_DIR, 'create_venv'),
+ _VENV_DIR,
+ _REQUIREMENTS,
+ ], stdout=sys.stderr)
+
+
+def _ExecInVenv(args):
+ """Exec command in chromite venv.
+
+ Args:
+ args: Sequence of arguments.
+ """
+ os.execv(os.path.join(_VIRTUALENV_DIR, 'venv_command'),
+ ['venv_command', _VENV_DIR] + list(args))
+
+
+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', '')
if __name__ == '__main__':
- if _IN_VENV in os.environ:
- wrapper.DoMain()
- else:
- create_cmd = os.path.join(_CHROMITE_DIR, 'venv', 'create_env.sh')
- cros_build_lib.RunCommand([create_cmd])
- python_cmd = os.path.join(_CHROMITE_DIR, 'venv', 'venv', 'bin', 'python')
- cmd = [python_cmd] + sys.argv
- o = cros_build_lib.RunCommand(
- cmd, extra_env={_IN_VENV: '1'},
- mute_output=False, error_code_ok=True)
- exit(o.returncode)
+ main()