pw_env_setup: Always set dir_pigweed for venv

Always set dir_pigweed for Python virtualenv setup.

Change-Id: I8c0712fe739ce518b899171a9a96e876d67a0407
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/33125
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>
Pigweed-Auto-Submit: Rob Mohr <mohrr@google.com>
Reviewed-by: Wyatt Hepler <hepler@google.com>
diff --git a/pw_env_setup/py/pw_env_setup/virtualenv_setup/install.py b/pw_env_setup/py/pw_env_setup/virtualenv_setup/install.py
index 41d545b..8e798ea 100644
--- a/pw_env_setup/py/pw_env_setup/virtualenv_setup/install.py
+++ b/pw_env_setup/py/pw_env_setup/virtualenv_setup/install.py
@@ -162,6 +162,8 @@
     venv_python = os.path.join(venv_bin, 'python')
 
     pw_root = os.environ.get('PW_ROOT')
+    if not pw_root and env:
+        pw_root = env.PW_ROOT
     if not pw_root:
         pw_root = git_repo_root()
     if not pw_root:
@@ -210,7 +212,13 @@
         gn_log_path = os.path.join(venv_path, gn_log)
         try:
             with open(gn_log_path, 'w') as outs:
-                subprocess.check_call(('gn', 'gen', build_dir),
+                gn_cmd = (
+                    'gn',
+                    'gen',
+                    build_dir,
+                    '--args=dir_pigweed="{}"'.format(pw_root),
+                )
+                subprocess.check_call(gn_cmd,
                                       cwd=os.path.join(project_root,
                                                        gn_target.directory),
                                       stdout=outs,