1// Copyright 2023 The Bazel Authors. All rights reserved. 2// 3// Licensed under the Apache License, Version 2.0 (the "License"); 4// you may not use this file except in compliance with the License. 5// You may obtain a copy of the License at 6// 7// http://www.apache.org/licenses/LICENSE-2.0 8// 9// Unless required by applicable law or agreed to in writing, software 10// distributed under the License is distributed on an "AS IS" BASIS, 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12// See the License for the specific language governing permissions and 13// limitations under the License. 14 15package python 16 17import ( 18 "bufio" 19 "context" 20 "fmt" 21 "io" 22 "log" 23 "os" 24 "os/exec" 25 "strconv" 26 "strings" 27 "sync" 28 29 "github.com/bazelbuild/rules_go/go/tools/bazel" 30) 31 32var ( 33 stdModulesStdin io.WriteCloser 34 stdModulesStdout io.Reader 35 stdModulesMutex sync.Mutex 36 stdModulesSeen map[string]struct{} 37) 38 39func startStdModuleProcess(ctx context.Context) { 40 stdModulesSeen = make(map[string]struct{}) 41 42 stdModulesScriptRunfile, err := bazel.Runfile("python/std_modules") 43 if err != nil { 44 log.Printf("failed to initialize std_modules: %v\n", err) 45 os.Exit(1) 46 } 47 48 cmd := exec.CommandContext(ctx, stdModulesScriptRunfile) 49 50 cmd.Stderr = os.Stderr 51 // All userland site-packages should be ignored. 52 cmd.Env = []string{"PYTHONNOUSERSITE=1"} 53 stdin, err := cmd.StdinPipe() 54 if err != nil { 55 log.Printf("failed to initialize std_modules: %v\n", err) 56 os.Exit(1) 57 } 58 stdModulesStdin = stdin 59 60 stdout, err := cmd.StdoutPipe() 61 if err != nil { 62 log.Printf("failed to initialize std_modules: %v\n", err) 63 os.Exit(1) 64 } 65 stdModulesStdout = stdout 66 67 if err := cmd.Start(); err != nil { 68 log.Printf("failed to initialize std_modules: %v\n", err) 69 os.Exit(1) 70 } 71 72 go func() { 73 if err := cmd.Wait(); err != nil { 74 log.Printf("failed to wait for std_modules: %v\n", err) 75 os.Exit(1) 76 } 77 }() 78} 79 80func shutdownStdModuleProcess() { 81 if err := stdModulesStdin.Close(); err != nil { 82 fmt.Fprintf(os.Stderr, "error closing std module: %v", err) 83 } 84} 85 86func isStdModule(m module) (bool, error) { 87 if _, seen := stdModulesSeen[m.Name]; seen { 88 return true, nil 89 } 90 stdModulesMutex.Lock() 91 defer stdModulesMutex.Unlock() 92 93 fmt.Fprintf(stdModulesStdin, "%s\n", m.Name) 94 95 stdoutReader := bufio.NewReader(stdModulesStdout) 96 line, err := stdoutReader.ReadString('\n') 97 if err != nil { 98 return false, err 99 } 100 if len(line) == 0 { 101 return false, fmt.Errorf("unexpected empty output from std_modules") 102 } 103 104 isStd, err := strconv.ParseBool(strings.TrimSpace(line)) 105 if err != nil { 106 return false, err 107 } 108 109 if isStd { 110 stdModulesSeen[m.Name] = struct{}{} 111 return true, nil 112 } 113 return false, nil 114} 115