kitchen: add -python-path

client/cmd/kitchen/cook.go

 // Copyright 2016 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 main
 
 import (
         "encoding/json"
         "fmt"
         "io/ioutil"
         "net/url"
         "os"
         "os/exec"
         "path"
+        "path/filepath"
         "strconv"
         "strings"
         "syscall"
 
         "github.com/maruel/subcommands"
         "golang.org/x/net/context"
 
         "github.com/luci/luci-go/common/cli"
         "github.com/luci/luci-go/common/clock"
         "github.com/luci/luci-go/common/ctxcmd"
+        "github.com/luci/luci-go/common/flag/stringlistflag"
 )
 
 // BootstrapStepName is the name of kitchen's step where it makes preparations
 // for running a recipe, e.g. fetches a repository.
 const BootstrapStepName = "recipe bootstrap"
 
 // cmdCook checks out a repository at a revision and runs a recipe.
 var cmdCook = &subcommands.Command{
         UsageLine: "cook -repository <repository URL> -revision <revision> -recipe <recipe>",
         ShortDesc: "Checks out a repository and runs a recipe.",
         LongDesc:  "Clones or fetches a repository, checks out a revision and runs a recipe",
         CommandRun: func() subcommands.CommandRun {
                 var c cookRun
                 fs := &c.Flags
                 fs.StringVar(&c.RepositoryURL, "repository", "", "URL of a git repository to fetch")
                 fs.StringVar(
                         &c.Revision,
                         "revision",
                         "FETCH_HEAD",
                         "Git commit hash to check out.")
                 fs.StringVar(&c.Recipe, "recipe", "<recipe>", "Name of the recipe to run")
+                fs.Var(&c.PythonPaths, "python-path", "Python path to include. Can be specified multiple times.")
                 fs.StringVar(
                         &c.CheckoutDir,
                         "checkout-dir",
                         "",
                         "The directory to check out the repository to. "+
                                 "Defaults to ./<repo name>, where <repo name> is the last component of -repository.")
                 fs.StringVar(
                         &c.Workdir,
                         "workdir",
                         "",
@@ skipping 21 matching lines @@
 
         RepositoryURL        string
         Revision             string
         Recipe               string
         CheckoutDir          string
         Workdir              string
         Properties           string
         PropertiesFile       string
         OutputResultJSONFile string
         Timestamps           bool
+        PythonPaths          stringlistflag.Flag
 }
 
 func (c *cookRun) validateFlags() error {
         // Validate Repository.
         if c.RepositoryURL == "" {
                 return fmt.Errorf("-repository is required")
         }
         repoURL, err := url.Parse(c.RepositoryURL)
         if err != nil {
                 return fmt.Errorf("invalid repository %q: %s", repoURL, err)
@@ skipping 27 matching lines @@
 
         if c.Workdir == "" {
                 var tempWorkdir string
                 if tempWorkdir, err = ioutil.TempDir("", "kitchen-"); err != nil {
                         return 0, err
                 }
                 defer os.RemoveAll(tempWorkdir)
                 c.Workdir = tempWorkdir
         }
 
+        for i, p := range c.PythonPaths {
+                p := filepath.FromSlash(p)
+                p, err := filepath.Abs(p)
+                if err != nil {
+                        return 0, err
+                }
+                c.PythonPaths[i] = p
+        }
+        // Why here? It is much easier than manipulating exec.Command.Env.
+        os.Setenv("PYTHONPATH", strings.Join(c.PythonPaths, string(os.PathListSeparator)))
+
         recipe := recipeRun{
                 repositoryPath:       c.CheckoutDir,
                 workDir:              c.Workdir,
                 recipe:               c.Recipe,
                 propertiesJSON:       c.Properties,
                 propertiesFile:       c.PropertiesFile,
                 outputResultJSONFile: c.OutputResultJSONFile,
                 timestamps:           c.Timestamps,
         }
         recipeCmd, err := recipe.Command()
@@ skipping 109 matching lines @@
 }
 
 func annotateTime(ctx context.Context) {
         timestamp := clock.Get(ctx).Now().Unix()
         annotate("CURRENT_TIMESTAMP", strconv.FormatInt(timestamp, 10))
 }
 
 func annotate(args ...string) {
         fmt.Printf("@@@%s@@@\n", strings.Join(args, "@"))
 }