lib/cmd/cmd.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 // package cmd defines a RunFunc type, representing a process that can
   6 // be invoked from a command line.
   7 package cmd
   8
   9 import (
  10         "flag"
  11         "fmt"
  12         "io"
  13 )
  14
  15 // A RunFunc runs a command with the given args, and returns an exit
  16 // code.
  17 type RunFunc func(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int
  18
  19 // Multi returns a command that looks up its first argument in m, and
  20 // runs the resulting RunFunc with the remaining args.
  21 //
  22 // Example:
  23 //
  24 //     os.Exit(Multi(map[string]RunFunc{
  25 //             "foobar": func(prog string, args []string) int {
  26 //                     fmt.Println(args[0])
  27 //                     return 2
  28 //             },
  29 //     })("/usr/bin/multi", []string{"foobar", "baz"}))
  30 //
  31 // ...prints "baz" and exits 2.
  32 func Multi(m map[string]RunFunc) RunFunc {
  33         return func(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  34                 if len(args) < 1 {
  35                         fmt.Fprintf(stderr, "usage: %s command [args]", prog)
  36                         return 2
  37                 }
  38                 if cmd, ok := m[args[0]]; !ok {
  39                         fmt.Fprintf(stderr, "unrecognized command %q", args[0])
  40                         return 2
  41                 } else {
  42                         return cmd(prog+" "+args[0], args[1:], stdin, stdout, stderr)
  43                 }
  44         }
  45 }
  46
  47 // WithLateSubcommand wraps a RunFunc by skipping over some known
  48 // flags to find a subcommand, and moving that subcommand to the front
  49 // of the args before calling the wrapped RunFunc. For example:
  50 //
  51 //      // Translate [           --format foo subcommand bar]
  52 //      //        to [subcommand --format foo            bar]
  53 //      WithLateSubcommand(fn, []string{"format"}, nil)
  54 func WithLateSubcommand(run RunFunc, argFlags, boolFlags []string) RunFunc {
  55         return func(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  56                 flags := flag.NewFlagSet("prog", flag.ContinueOnError)
  57                 for _, arg := range argFlags {
  58                         flags.String(arg, "", "")
  59                 }
  60                 for _, arg := range boolFlags {
  61                         flags.Bool(arg, false, "")
  62                 }
  63                 // Ignore errors. We can't report a useful error
  64                 // message anyway.
  65                 flags.Parse(args)
  66                 if flags.NArg() > 0 {
  67                         // Move the first arg after the recognized
  68                         // flags up to the front.
  69                         flagargs := len(args) - flags.NArg()
  70                         newargs := make([]string, len(args))
  71                         newargs[0] = args[flagargs]
  72                         copy(newargs[1:flagargs+1], args[:flagargs])
  73                         copy(newargs[flagargs+1:], args[flagargs+1:])
  74                         args = newargs
  75                 }
  76                 return run(prog, args, stdin, stdout, stderr)
  77         }
  78 }