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()