scripts: add virtualenv_wrapper

BUG=chromium:645611
TEST=~/chromiumos/chromite/bin$ ./venv_check -> succeeds and prints the
correct import path. Also succeeds when run from other cwds.

Change-Id: I981ed7d25457381688915171c04a249457a219b5
Reviewed-on: https://chromium-review.googlesource.com/388773
Commit-Ready: Aviv Keshet <akeshet@chromium.org>
Tested-by: Aviv Keshet <akeshet@chromium.org>
Reviewed-by: Paul Hobbs <phobbs@google.com>
diff --git a/scripts/virtualenv_wrapper.py b/scripts/virtualenv_wrapper.py
new file mode 100755
index 0000000..4852325
--- /dev/null
+++ b/scripts/virtualenv_wrapper.py
@@ -0,0 +1,44 @@
+#!/usr/bin/env python2
+# Copyright 2016 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Wrapper around chromite executable scripts that use virtualenv."""
+
+from __future__ import print_function
+
+import os
+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
+
+# 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.cbuildbot import constants
+
+_CHROMITE_DIR = constants.CHROMITE_DIR
+_IN_VENV = 'IN_CHROMITE_VENV'
+
+
+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)