pw_env_setup: Add --virtualenv-root option

Add --virtualenv-root option to pw_env_setup. This makes it easier to
use pw_env_setup from LUCI.

Change-Id: I6f1b4003f0f56acc412bfed2e0ff9fa70504c77c
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/22800
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>
Reviewed-by: David Rogers <davidrogers@google.com>
diff --git a/pw_env_setup/py/pw_env_setup/env_setup.py b/pw_env_setup/py/pw_env_setup/env_setup.py
index b129065..244194a 100755
--- a/pw_env_setup/py/pw_env_setup/env_setup.py
+++ b/pw_env_setup/py/pw_env_setup/env_setup.py
@@ -172,7 +172,7 @@
 class EnvSetup(object):
     """Run environment setup for Pigweed."""
     def __init__(self, pw_root, cipd_cache_dir, shell_file, quiet, install_dir,
-                 use_pigweed_defaults, cipd_package_file,
+                 use_pigweed_defaults, cipd_package_file, virtualenv_root,
                  virtualenv_requirements, virtualenv_setup_py_root,
                  cargo_package_file, enable_cargo, json_file, *args, **kwargs):
         super(EnvSetup, self).__init__(*args, **kwargs)
@@ -185,6 +185,8 @@
         self._is_windows = os.name == 'nt'
         self._quiet = quiet
         self._install_dir = install_dir
+        self._virtualenv_root = (virtualenv_root
+                                 or os.path.join(install_dir, 'pigweed-venv'))
 
         if os.path.isfile(shell_file):
             os.unlink(shell_file)
@@ -373,8 +375,6 @@
     def virtualenv(self):
         """Setup virtualenv."""
 
-        venv_path = os.path.join(self._install_dir, 'python3-env')
-
         requirements, req_glob_warnings = _process_globs(
             self._virtualenv_requirements)
         setup_py_roots, setup_glob_warnings = _process_globs(
@@ -400,7 +400,7 @@
         if not requirements and not setup_py_roots:
             return result(_Result.Status.SKIPPED)
 
-        if not virtualenv_setup.install(venv_path=venv_path,
+        if not virtualenv_setup.install(venv_path=self._virtualenv_root,
                                         requirements=requirements,
                                         setup_py_roots=setup_py_roots,
                                         python=new_python3,
@@ -508,6 +508,13 @@
     )
 
     parser.add_argument(
+        '--virtualenv-root',
+        help=('Basename of virtualenv directory. Default: '
+              '<install_dir>/pigweed-venv'),
+        default=None,
+    )
+
+    parser.add_argument(
         '--cargo-package-file',
         help='Rust cargo packages to install. Lines with package name and '
         'version separated by a space.',