[vpython] Implement smart probing.

vpython/application/application.go

 // Copyright 2017 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 package application
 
 import (
         "flag"
         "fmt"
         "io/ioutil"
@@ skipping 12 matching lines @@
         "github.com/luci/luci-go/vpython/venv"
 
         cipdVersion "github.com/luci/luci-go/cipd/version"
         "github.com/luci/luci-go/common/cli"
         "github.com/luci/luci-go/common/errors"
         "github.com/luci/luci-go/common/logging"
         "github.com/luci/luci-go/common/logging/gologger"
         "github.com/luci/luci-go/common/system/environ"
         "github.com/luci/luci-go/common/system/exitcode"
         "github.com/luci/luci-go/common/system/filesystem"
+        "github.com/luci/luci-go/common/system/prober"
 )
 
 const (
         // VirtualEnvRootENV is an environment variable that, if set, will be used
         // as the default VirtualEnv root.
         //
         // This value overrides the default (~/.vpython), but can be overridden by the
         // "-root" flag.
         //
         // Like "-root", if this value is present but empty, a tempdir will be used
@@ skipping 15 matching lines @@
 type VerificationFunc func(context.Context, string, venv.PackageLoader, *vpythonAPI.Environment)
 
 // Config is an application's default configuration.
 type Config struct {
         // PackageLoader is the package loader to use.
         PackageLoader venv.PackageLoader
 
         // VENVPackage is the VirtualEnv package to use for bootstrap generation.
         VENVPackage vpythonAPI.Spec_Package
 
+        // RelativePathOverride is a series of forward-slash-delimited paths to
+        // directories relative to the "vpython" executable that will be checked
+        // for Python targets prior to checking PATH. This allows bundles (e.g., CIPD)
+        // that include both the wrapper and a real implementation, to force the
+        // wrapper to use the bundled implementation if present.
+        //
+        // See "github.com/luci/luci-go/common/wrapper/prober.Probe"'s
+        // RelativePathOverride member for more information.
+        RelativePathOverride []string
+
         // PruneThreshold, if > 0, is the maximum age of a VirtualEnv before it
         // becomes candidate for pruning. If <= 0, no pruning will be performed.
         //
         // See venv.Config's PruneThreshold.
         PruneThreshold time.Duration
         // MaxPrunesPerSweep, if > 0, is the maximum number of VirtualEnv that should
         // be pruned passively. If <= 0, no limit will be applied.
         //
         // See venv.Config's MaxPrunesPerSweep.
         MaxPrunesPerSweep int
@@ skipping 60 matching lines @@
         fs.StringVar(&a.opts.EnvConfig.BaseDir, "root", a.opts.EnvConfig.BaseDir,
                 "Path to virtual environment root directory. Default is the working directory. "+
                         "If explicitly set to empty string, a temporary directory will be used and cleaned up "+
                         "on completion.")
         fs.StringVar(&a.specPath, "spec", a.specPath,
                 "Path to environment specification file to load. Default probes for one.")
 
         a.logConfig.AddFlags(fs)
 }
 
-func (a *application) mainImpl(c context.Context, args []string) error {
+func (a *application) mainImpl(c context.Context, argv0 string, args []string) error {
         // Determine our VirtualEnv base directory.
         if v, ok := a.opts.Environ.Get(VirtualEnvRootENV); ok {
                 a.opts.EnvConfig.BaseDir = v
         } else {
                 hdir, err := homedir.Dir()
                 if err != nil {
                         return errors.Annotate(err).Reason("failed to get user home directory").Err()
                 }
                 a.opts.EnvConfig.BaseDir = filepath.Join(hdir, ".vpython")
         }
 
         // Extract "vpython" arguments and parse them.
         fs := flag.NewFlagSet("", flag.ExitOnError)
         fs.SetOutput(os.Stdout) // Python uses STDOUT for help and flag information.
 
         a.addToFlagSet(fs)
         selfArgs, args := extractFlagsForSet(args, fs)
         if err := fs.Parse(selfArgs); err != nil && err != flag.ErrHelp {
                 return errors.Annotate(err).Reason("failed to parse flags").Err()
         }
 
+        // Identify the "self" executable. Use this to construct a "lookPath", which
+        // will be used to locate the base Python interpreter.
+        lp := lookPath{
+                probeBase: prober.Probe{
+                        RelativePathOverride: a.RelativePathOverride,
+                },
+                env: a.opts.Environ,
+        }
+        if err := lp.probeBase.ResolveSelf(argv0); err != nil {
+                logging.WithError(err).Warningf(c, "Failed to resolve 'self'")
+        }
+        a.opts.EnvConfig.LookPathFunc = lp.look
+
         if a.help {
-                return a.showPythonHelp(c, fs)
+                return a.showPythonHelp(c, fs, &lp)
         }
 
         c = a.logConfig.Set(c)
 
         // If an spec path was manually specified, load and use it.
         if a.specPath != "" {
                 var err error
                 if a.opts.EnvConfig.Spec, err = spec.Load(a.specPath); err != nil {
                         return errors.Annotate(err).Reason("failed to load specification file (-spec) from: %(path)s").
                                 D("path", a.specPath).
@@ skipping 28 matching lines @@
                 // as our error.
                 if rc, has := exitcode.Get(errors.Unwrap(err)); has {
                         err = ReturnCodeError(rc)
                 }
 
                 return errors.Annotate(err).Err()
         }
         return nil
 }
 
-func (a *application) showPythonHelp(c context.Context, fs *flag.FlagSet) error {
+func (a *application) showPythonHelp(c context.Context, fs *flag.FlagSet, lp *lookPath) error {
         self, err := os.Executable()
         if err != nil {
                 self = "vpython"
         }
         vers, err := cipdVersion.GetStartupVersion()
         if err == nil && vers.PackageName != "" && vers.InstanceID != "" {
                 self = fmt.Sprintf("%s (%s@%s)", self, vers.PackageName, vers.InstanceID)
         }
 
         fmt.Fprintf(os.Stdout, "Usage of %s:\n", self)
         fs.SetOutput(os.Stdout)
         fs.PrintDefaults()
 
-        i, err := python.Find(c, python.Version{})
+        i, err := python.Find(c, python.Version{}, lp.look)
         if err != nil {
                 return errors.Annotate(err).Reason("could not find Python interpreter for help").Err()
         }
 
         // Redirect all "--help" to Stdout for consistency.
         fmt.Fprintf(os.Stdout, "\nPython help for %s:\n", i.Python)
         cmd := i.IsolatedCommand(c, "--help")
         cmd.Stdout = os.Stdout
         cmd.Stderr = os.Stdout
         if err := cmd.Run(); err != nil {
                 return errors.Annotate(err).Reason("failed to dump Python help from: %(interpreter)s").
                         D("interpreter", i.Python).
                         Err()
         }
         return nil
 }
 
 // Main is the main application entry point.
 func (cfg *Config) Main(c context.Context) int {
+        // Implementation of "checkWrapper": if CheckWrapperENV is set, we immediately
+        // exit with a non-zero value.
+        env := environ.System()
+        if wrapperCheck(env) {
+                return 1
+        }
+
         c = gologger.StdConfig.Use(c)
         c = logging.SetLevel(c, logging.Error)
 
         a := application{
                 Config: cfg,
                 opts: vpython.Options{
                         EnvConfig: venv.Config{
                                 BaseDir:           "", // (Determined below).
                                 MaxHashLen:        6,
                                 Package:           cfg.VENVPackage,
                                 PruneThreshold:    cfg.PruneThreshold,
                                 MaxPrunesPerSweep: cfg.MaxPrunesPerSweep,
                                 MaxScriptPathLen:  cfg.MaxScriptPathLen,
                                 Loader:            cfg.PackageLoader,
                         },
                         WaitForEnv: true,
-                        Environ:    environ.System(),
+                        Environ:    env,
                 },
                 logConfig: logging.Config{
                         Level: logging.Error,
                 },
         }
 
         return run(c, func(c context.Context) error {
-                return a.mainImpl(c, os.Args[1:])
+                return a.mainImpl(c, os.Args[0], os.Args[1:])
         })
 }