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.

2 files changed, 15 insertions, 17 deletions

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",