gopkg.in/hugelgupf/u-root.v2@v2.0.0-20180831055005-3f8fdb0ce09d/pkg/bb/register.go (about)

     1  // Copyright 2018 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package bb
     6  
     7  import (
     8  	"errors"
     9  	"fmt"
    10  	"os"
    11  )
    12  
    13  // ErrNotRegistered is returned by Run if the given command is not registered.
    14  var ErrNotRegistered = errors.New("command not registered")
    15  
    16  // Noop is a noop function.
    17  var Noop = func() {}
    18  
    19  type bbCmd struct {
    20  	init, main func()
    21  }
    22  
    23  var bbCmds = map[string]bbCmd{}
    24  
    25  var defaultCmd *bbCmd
    26  
    27  // Register registers an init and main function for name.
    28  func Register(name string, init, main func()) {
    29  	if _, ok := bbCmds[name]; ok {
    30  		panic(fmt.Sprintf("cannot register two commands with name %q", name))
    31  	}
    32  	bbCmds[name] = bbCmd{
    33  		init: init,
    34  		main: main,
    35  	}
    36  }
    37  
    38  // RegisterDefault registers a default init and main function.
    39  func RegisterDefault(init, main func()) {
    40  	defaultCmd = &bbCmd{
    41  		init: init,
    42  		main: main,
    43  	}
    44  }
    45  
    46  // Run runs the command with the given name.
    47  //
    48  // If the command's main exits without calling os.Exit, Run will exit with exit
    49  // code 0.
    50  func Run(name string) error {
    51  	var cmd *bbCmd
    52  	if c, ok := bbCmds[name]; ok {
    53  		cmd = &c
    54  	} else if defaultCmd != nil {
    55  		cmd = defaultCmd
    56  	} else {
    57  		return ErrNotRegistered
    58  	}
    59  	cmd.init()
    60  	cmd.main()
    61  	os.Exit(0)
    62  	// Unreachable.
    63  	return nil
    64  }