Evolution of main.go file
I happen to write a lot of command line tools in Go. Below are few useful strategies I came up over time.
Handling errors
Lots of examples of main() function follow the same pattern:
func main() {
...
if err := nil { log.Fatal(err) }
...
if err := nil { log.Fatal(err) }
...
}
That is, main since main function does not return anything, people handle errors all over this function and call log.Fatal(…) to terminate program.
This approach is suboptimal for the following reasons:
- It does not look clean. main function code is riddled with ad-hoc error handling. If you'd ever need to refactor this function to make it standalone, all of these
log.Fatal(…) calls would likely need to be changed to something else.
- Deferred functions cannot be relied upon in main function when considering "non-success" execution path. This is because of
log.Fatal(…), which is documented to be an equivalent of log.Print() followed by a call to os.Exit(1), that itself states that program terminates immediately; deferred functions are not run.
Luckily, this is easily fixable by restructuring the program as follows:
func main() {
if err := run() ; err !=nil {
log.Fatal(err)
}
}
func run() error {
...
if err !=nil { return err }
...
return err
}
Notice that now all main logic lives inside run() function which returns error, and main() function now looks really slim and only handles single error comming from run(). This has a nice property of being able to use defers inside run(), this function may also be reused almost without any refactoring if this package ever gets promoted to a library.
Dealing with flags
Global state should be avoided, so any variables should be explicitly passed to function. This also applies to flags. Consider above example of main() function being a slim wrapper around run() that contains all the logic. If program requires configuration via flags, what would be the best place to put them? Obviously, to keep run() being easy to decople if needed, we need to define flags outside it and only pass it defined values.
Consider this example:
func main() {
var dir string
var rec bool
flag.StringVar(&dir, "dir", "", "path to directory")
flag.BoolVar(&rec, "r", false, "traverse directory recursively")
flag.Parse()
if err := run(dir, rec) ; err !=nil {
log.Fatal(err)
}
}
func run(dir string, rec bool) error {
...
return err
}
This works quite nicely, but what if we're going to add another flag? And then another? run() function signature would then grow more and more. There's a better way to do so — group flag-related arguments under the same struct type:
func main() {
var args runArgs
flag.StringVar(&args.Dir, "dir", "", "path to directory")
flag.BoolVar(&args.Rec, "r", false, "traverse directory recursively")
flag.Parse()
if err := run(args) ; err !=nil {
log.Fatal(err)
}
}
type runArgs struct {
Dir string
Rec bool
}
func run(args runArgs) error {
...
return err
}
With such dedicated type it is easier to add another argument/flag to our run() function. After using this style for a while I found it to be almost what I wanted, and then decided to take it a bit further, by automating all flag definitions (those flag.*Var calls) introducing autoflags package which is a small wrapper on top of flag from the standard library.
With autoflags package the above example transforms into this:
func main() {
args := runArgs{
Dir: "", // default values can be defined here
}
autoflags.Parse(&args)
if err := run(args) ; err !=nil {
log.Fatal(err)
}
}
type runArgs struct {
Dir string `flag:"dir,path to directory"`
Rec bool `flag:"r,traverse directory recursively"`
}
func run(args runArgs) error {
...
return err
}
Now adding yet another flag requires only adding extra field to runArgs type.
It is a matter of good manners to document Go code, even if it's a command line tool.
But what if it was possible to reuse main package documentation as program output shown when run with -h (help) flag?
Below is the code I came up with to do exactly this.
//+build generate
package main
import (
"bytes"
"fmt"
"go/doc"
"go/parser"
"go/token"
"io/ioutil"
"os"
)
func main() {
if err := run(); err != nil {
os.Stderr.WriteString(err.Error() + "\n")
os.Exit(1)
}
}
func run() error {
fset := token.NewFileSet()
m, err := parser.ParseDir(fset, ".", nil, parser.ParseComments)
if err != nil {
return err
}
p, ok := m["main"]
if !ok {
return fmt.Errorf("cannot find main package")
}
docBuf := new(bytes.Buffer)
for _, f := range p.Files {
if f.Doc == nil {
continue
}
doc.ToText(docBuf, f.Doc.Text(), "", "\t", 80)
}
if docBuf.Len() == 0 {
return fmt.Errorf("could not extract any docs")
}
buf := new(bytes.Buffer)
fmt.Fprintf(buf, templateFormat, docBuf.String())
return ioutil.WriteFile("usage_generated.go", buf.Bytes(), 0666)
}
const templateFormat = `// go run update_usage.go
// Code generated by the command above; DO NOT EDIT.
package main
const usage = %#v
`
Save that code under update_usage.go file in the same directory as your main package, then add this line somewhere inside your main package source:
//go:generate go run update_usage.go
Now, whenever you run go generate, it will recreate usage_generated.go file holding global usage constant, holding text of your package's documentation (what is shown by go doc). You can then include both update_usage.go and usage_generated.go into VCS along with other source files, just don't forget to run go generate whenever you update the package documentation.
Now to actually use this text as a help shown on -h flag, redefine flag.Usage variable like this somewhere in your code:
func init() {
flag.Usage = func() {
fmt.Fprintln(flag.CommandLine.Output(), usage)
fmt.Fprintf(flag.CommandLine.Output(), "Usage: %s [flags]\n", filepath.Base(os.Args[0]))
flag.PrintDefaults()
}
}
Now when your command is run with -h flag it will output package-level documentation followed by list of supported flags with their defaults and description.
Alternatively you can use a standalone usagegen tool that implements above logic.