go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/cli/cli.go (about)

     1  // Copyright 2016 The LUCI Authors.
     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  
    15  // Package cli is a helper package for "github.com/maruel/subcommands".
    16  //
    17  // It adds a non-intrusive integration with context.Context.
    18  package cli
    19  
    20  import (
    21  	"context"
    22  	"io"
    23  	"os"
    24  
    25  	"github.com/maruel/subcommands"
    26  
    27  	"go.chromium.org/luci/common/system/environ"
    28  )
    29  
    30  // ContextModificator takes a context, adds something, and returns a new one.
    31  //
    32  // It is implemented by Application and can optionally by implemented by
    33  // subcommands.CommandRun instances. It will be called by GetContext to modify
    34  // initial context.
    35  type ContextModificator interface {
    36  	ModifyContext(context.Context) context.Context
    37  }
    38  
    39  // GetContext sniffs ContextModificator in the app and in the cmd and uses them
    40  // to derive a context for the command.
    41  //
    42  // Subcommands can use it to get an initial context in their `Run` methods.
    43  //
    44  // Uses `env` to initialize luci/common/system/environ environment in the
    45  // context. In particular populates unset (but registered in the CLI app) env
    46  // vars with their default values (if they are not empty).
    47  //
    48  // Order of the context modifications:
    49  //  1. Start with the background context.
    50  //  2. Apply `app` modifications if `app` implements ContextModificator.
    51  //  3. Initialize luci/common/system/environ in the context.
    52  //  4. Apply `cmd` modifications if `cmd` implements ContextModificator.
    53  //
    54  // In particular, command's modificator is able to examine env vars in
    55  // the context.
    56  func GetContext(app subcommands.Application, cmd subcommands.CommandRun, env subcommands.Env) context.Context {
    57  	ctx := context.Background()
    58  	if m, _ := app.(ContextModificator); m != nil {
    59  		ctx = m.ModifyContext(ctx)
    60  	}
    61  
    62  	root := environ.FromCtx(ctx)
    63  	for name, val := range env {
    64  		if val.Value != "" || val.Exists {
    65  			root.Set(name, val.Value)
    66  		}
    67  	}
    68  	ctx = root.SetInCtx(ctx)
    69  
    70  	if m, _ := cmd.(ContextModificator); m != nil {
    71  		ctx = m.ModifyContext(ctx)
    72  	}
    73  
    74  	return ctx
    75  }
    76  
    77  // Application is like subcommands.DefaultApplication, except it also implements
    78  // ContextModificator.
    79  type Application struct {
    80  	Name     string
    81  	Title    string
    82  	Context  func(context.Context) context.Context
    83  	Commands []*subcommands.Command
    84  	EnvVars  map[string]subcommands.EnvVarDefinition
    85  
    86  	profiling profilingExt // empty struct if 'include_profiler' build tag is not set
    87  }
    88  
    89  var _ interface {
    90  	subcommands.Application
    91  	ContextModificator
    92  } = (*Application)(nil)
    93  
    94  // GetName implements interface subcommands.Application.
    95  func (a *Application) GetName() string {
    96  	return a.Name
    97  }
    98  
    99  // GetTitle implements interface subcommands.Application.
   100  func (a *Application) GetTitle() string {
   101  	return a.Title
   102  }
   103  
   104  // GetCommands implements interface subcommands.Application.
   105  func (a *Application) GetCommands() []*subcommands.Command {
   106  	a.profiling.addProfiling(a.Commands)
   107  	return a.Commands
   108  }
   109  
   110  // GetOut implements interface subcommands.Application.
   111  func (a *Application) GetOut() io.Writer {
   112  	return os.Stdout
   113  }
   114  
   115  // GetErr implements interface subcommands.Application.
   116  func (a *Application) GetErr() io.Writer {
   117  	return os.Stderr
   118  }
   119  
   120  // GetEnvVars implements interface subcommands.Application.
   121  func (a *Application) GetEnvVars() map[string]subcommands.EnvVarDefinition {
   122  	return a.EnvVars
   123  }
   124  
   125  // ModifyContext implements interface ContextModificator.
   126  //
   127  // 'ctx' here is always context.Background().
   128  func (a *Application) ModifyContext(ctx context.Context) context.Context {
   129  	if a.Context != nil {
   130  		return a.Context(ctx)
   131  	}
   132  	return ctx
   133  }