From 767e3ced63df806242c3f5103b1290bc3ce4b914 Mon Sep 17 00:00:00 2001
From: aptenodytes-forsteri
 <92043606+aptenodytes-forsteri@users.noreply.github.com>
Date: Mon, 23 Jan 2023 16:58:33 -0500
Subject: Redirect stdout when checking imports. (#1007)

Fixes https://github.com/bazelbuild/rules_python/issues/1006

Also, set PYTHONNOUSERSITE so that the script doesn't even look in site
packages when checking modules.

Fix typo with capitilize.
---
 gazelle/std_modules.go |  4 ++--
 gazelle/std_modules.py | 28 +++++++++++++---------------
 2 files changed, 15 insertions(+), 17 deletions(-)

(limited to 'gazelle')

diff --git a/gazelle/std_modules.go b/gazelle/std_modules.go
index f7d0c24..e784a2d 100644
--- a/gazelle/std_modules.go
+++ b/gazelle/std_modules.go
@@ -37,8 +37,8 @@ func init() {
 	cmd := exec.CommandContext(ctx, stdModulesScriptRunfile)
 
 	cmd.Stderr = os.Stderr
-	cmd.Env = []string{}
-
+	// All userland site-packages should be ignored.
+	cmd.Env = []string{"PYTHONNOUSERSITE=1"}
 	stdin, err := cmd.StdinPipe()
 	if err != nil {
 		log.Printf("failed to initialize std_modules: %v\n", err)
diff --git a/gazelle/std_modules.py b/gazelle/std_modules.py
index ccd1dcd..86a2077 100644
--- a/gazelle/std_modules.py
+++ b/gazelle/std_modules.py
@@ -3,30 +3,28 @@
 # it evaluates, it outputs true/false for whether the module is part of the
 # standard library or not.
 
-import site
+import os
 import sys
-
-
-# Don't return any paths, all userland site-packages should be ignored.
-def __override_getusersitepackages__():
-    return ""
-
-
-site.getusersitepackages = __override_getusersitepackages__
+from contextlib import redirect_stdout
 
 
 def is_std_modules(module):
-    try:
-        __import__(module, globals(), locals(), [], 0)
-        return True
-    except Exception:
-        return False
+    # If for some reason a module (such as pygame, see https://github.com/pygame/pygame/issues/542)
+    # prints to stdout upon import,
+    # the output of this script should still be parseable by golang.
+    # Therefore, redirect stdout while running the import.
+    with redirect_stdout(os.devnull):
+        try:
+            __import__(module, globals(), locals(), [], 0)
+            return True
+        except Exception:
+            return False
 
 
 def main(stdin, stdout):
     for module in stdin:
         module = module.strip()
-        # Don't print the boolean directly as it is captilized in Python.
+        # Don't print the boolean directly as it is capitalized in Python.
         print(
             "true" if is_std_modules(module) else "false",
             end="\n",
-- 
cgit v1.2.3