[vpython] Implement smart probing.

vpython/application/probe.go

Index: vpython/application/probe.go
diff --git a/vpython/application/probe.go b/vpython/application/probe.go
new file mode 100644
index 0000000000000000000000000000000000000000..a12713e5bc986c3c3a4e95fbcf29219fd6dfe1b0
--- /dev/null
+++ b/vpython/application/probe.go
@@ -0,0 +1,137 @@
+// 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 (
+	"fmt"
+	"os"
+	"os/exec"
+	"strings"
+
+	"golang.org/x/net/context"
+
+	"github.com/luci/luci-go/vpython/python"
+
+	"github.com/luci/luci-go/common/errors"
+	"github.com/luci/luci-go/common/logging"
+	"github.com/luci/luci-go/common/system/environ"
+	"github.com/luci/luci-go/common/system/exitcode"
+	"github.com/luci/luci-go/common/system/prober"
+)
+
+const (
+	// CheckWrapperENV is an environment variable that is used to detect if a
+	// given binary is a "vpython" wrapper instance. The "vpython" Main function
+	// will immediately exit with a non-zero return code if this environment
+	// variable is set.
+	//
+	// See checkWrapper for more details.
+	checkWrapperENV = "VPYTHON_CHECK_WRAPPER"
+
+	// checkWrapperSentinel is text that will be output by a "vpython" instance if
+	// checkWrapperENV is set. If an application exits with a non-zero error code
+	// and outputs this text to STDOUT, then it is considered a "vpython" wrapper.
+	checkWrapperSentinel = checkWrapperENV + ": vpython wrapper check, exiting with non-zero return code"
+)
+
+func wrapperCheck(env environ.Env) bool {
+	if env.GetEmpty(checkWrapperENV) != "" {
+		fmt.Fprintln(os.Stdout, checkWrapperSentinel)
+		return true
+	}
+	return false
+}
+
+type lookPath struct {
+	// probeBase is a base prober whose Self and SelfStat fields will be used
+	// during a lookPathImpl call. This ensures that we only have to call
+	// ResolveSelf once for any given application.
+	probeBase prober.Probe
+
+	env        environ.Env
+	lastPath   string
+	lastOutput string
+}
+
+func (lp *lookPath) look(c context.Context, target string) (*python.LookPathResult, error) {
+	// Construct our probe, derived it from our "probeBase" and tailored to the
+	// specific "target" we're looking for.
+	probe := lp.probeBase
+	probe.Target = target
+	probe.CheckWrapper = lp.checkWrapper
+
+	var result python.LookPathResult
+	var err error
+	if result.Path, err = probe.Locate(c, "", lp.env.Clone()); err != nil {
+		return nil, err
+	}
+
+	// During probing, we try and capture the output of "--version". If we do,
+	// and if we confirm that the "path" returned by the probe matches "lastPath",
+	// return this version along with the resolution.
+	if result.Path == lp.lastPath {
+		var err error
+		if result.Version, err = python.ParseVersionOutput(lp.lastOutput); err == nil {
+			logging.Debugf(c, "Detected Python version %q from probe candidate [%s]", result.Version, result.Path)
+		} else {
+			logging.Debugf(c, "Failed to parse version from probe candidate [%s]: %s", result.Path, err)
+		}
+	}
+
+	return &result, nil
+}
+
+// checkWrapper is a prober.CheckWrapperFunc that detects if a given path is a
+// "vpython" instance.
+//
+// It does this by running it "path --version" with CheckWrapperENV set. If
+// the target is a "vpython" wrapper, it will immediately exit with a non-zero
+// value (see Main). If it is not a valid Python program, its behavior is
+// undefined. If it is a valid Pytohn instance, it will exit with a Python
+// version string.
+//
+// As a side effect, the text output of the last version probe will be stored
+// in lp.lastOutput. The "look" function can pass this value on to
+// the LookPathResult.
+func (lp *lookPath) checkWrapper(c context.Context, path string, env environ.Env) (isWrapper bool, err error) {
+	env.Set(checkWrapperENV, "1")
+
+	cmd := exec.CommandContext(c, path, "--version")
+	cmd.Env = env.Sorted()
+
+	output, err := cmd.CombinedOutput()
+	rc, ok := exitcode.Get(err)
+	if !ok {
+		err = errors.Annotate(err).Reason("failed to check if %(path)q is a wrapper").
+			D("path", path).
+			Err()
+		return
+	}
+
+	// Retain the last output, so our "look" can reference it to extract the
+	// version string.
+	lp.lastPath = path
+	lp.lastOutput = string(output)
+
+	if rc != 0 {
+		// The target exited with a non-zero error code. It is a wrapper if it
+		// printed the sentinel text.
+		isWrapper = strings.TrimSpace(lp.lastOutput) == checkWrapperSentinel
+		if isWrapper {
+			// The target was successfully identified as a wrapper. Clear "err", which
+			// is non-nil b/c of the non-zero exit code, to indicate a successful
+			// wrapper check.
+			err = nil
+			return
+		}
+
+		// The target returned non-zero, but didn't identify as a wrapper. It is
+		// likely something that happens to be named the same thing as the target,
+		// which is an error.
+		err = errors.Annotate(err).Reason("wrapper check returned non-zero").Err()
+	}
+
+	return
+}