Custom command-line flags with flag.Func

One of my favorite things about the recent Go 1.16 release is a small — but very welcome — addition to the flag package: the flag.Func() function. This makes it much easier to define and use custom command-line flags in your application.

For example, if you want to parse a flag like --pause=10s directly into a time.Duration type, or parse --urls="http://example.com http://example.org" directly into a []string slice, then previously you had two options. You could either create a custom type to implement the flag.Value interface, or use a third-party package like pflag.

But now the flag.Func() function gives you a simple and lightweight alternative. In this short post we're going to take a look at a few examples of how you can use it in your own code.

Parsing custom flag types

To demonstrate how this works, let's start with the two examples I gave above and create a sample application which accepts a list of URLs and then prints them out with a pause between them. Similar to this:

$ go run . --pause=3s --urls="http://example.com http://example.org http://example.net"
2021/03/08 08:16:04 http://example.com
2021/03/08 08:16:07 http://example.org
2021/03/08 08:16:10 http://example.net

To make this work, we'll need to do two things:

• Convert the --pause flag value from a 'human-readable' string like 200ms, 5s or 10m into a native Go time.Duration type. We can do this using the time.ParseDuration() function.
• Split the values in the --urls flag into a slice, so we can loop through them. The strings.Fields function is a good fit for this task.

We can use those together with flag.Func() like so:

package main

import (
    "flag"
    "log"
    "strings"
    "time"
)

func main() {
    // First we need to declare variables to hold the values from the
    // command-line flags. Notice that we also need to set any defaults,
    // which will be used if the relevant flag is not provided at runtime.
    var (
        urls  []string                    // Default of the empty slice
        pause time.Duration = time.Second // Default of one second
    )

    // The flag.Func() function takes three parameters: the flag name, 
    // descriptive help text, and a function with the signature 
    // `func(string) error` which is called to process the string value 
    // from the command-line flag at runtime and assign it to the necessary 
    // variable. In this case, we use strings.Fields() to split the string 
    // based on whitespace and store the resulting slice in the urls  
    // variable that we declared above. We then return nil from the 
    // function to indicate that the flag was parsed without any errors.
    flag.Func("urls", "List of URLs to print", func(flagValue string) error {
        urls = strings.Fields(flagValue)
        return nil
    })

    // Likewise we can do the same thing to parse the pause duration. The 
    // time.ParseDuration() function may throw an error here, so we make 
    // sure to return that from our function.
    flag.Func("pause", "Duration to pause between printing URLs", func(flagValue string) error {
        var err error
        pause, err = time.ParseDuration(flagValue)
        return err
    })

    // Importantly, call flag.Parse() to trigger actual parsing of the 
    // flags.
    flag.Parse()

    // Print out the URLs, pausing between each iteration.
    for _, u := range urls {
        log.Println(u)
        time.Sleep(pause)
    }
}

If you try to run this application, you should find that the flags are parsed and work just like you would expect. For example:

$ go run . --pause=500ms --urls="http://example.com http://example.org http://example.net"
2021/03/08 08:22:33 http://example.com
2021/03/08 08:22:34 http://example.org
2021/03/08 08:22:34 http://example.net

Whereas if you provide an invalid flag value that triggers an error in one of the flag.Func() functions, Go will automatically display the corresponding error message and exit. For example:

$ go run . --pause=500xx --urls="http://example.com http://example.org http://example.net"
invalid value "500xx" for flag -pause: time: unknown unit "xx" in duration "500xx"
Usage of /tmp/go-build3141872390/b001/exe/example.text:
  -pause value
        Duration to pause between printing URLs
  -urls value
        List of URLs to print
exit status 2

It's really important to point out here that if a flag isn't provided, the corresponding flag.Func() function will not be called at all. This means that you cannot set a default value inside a flag.Func() function, so trying to do something like this won't work:

flag.Func("pause", "Duration to pause between printing URLs (default 1s)", func(flagValue string) error {
    // DON'T DO THIS! This function wont' be called if the flag value is "".
    if flagValue == "" {
        pause = time.Second
        return nil
    }

    var err error
    pause, err = time.ParseDuration(flagValue)
    return err
})

On the plus side though, there are no restrictions on the code that can be contained in a flag.Func() function, so if you want, you could get even fancier with this and parse the URLs into a []*url.URL slice instead of a []string. Like so:

var (
    urls  []*url.URL                 
    pause time.Duration = time.Second
)

flag.Func("urls", "List of URLs to print", func(flagValue string) error {
    for _, u := range strings.Fields(flagValue) {
        parsedURL, err := url.Parse(u)
        if err != nil {
            return err
        }
        urls = append(urls, parsedURL)
    }
    return nil
})

Validating flag values

The flag.Func() function also opens up some new opportunities for validating input data from command-line flags. For example, let's say that your application has an --environment flag and you want to restrict the possible values to development, staging or production.

To do that, you can implement a flag.Func() function similar to this:

package main

import (
    "errors"
    "flag"
    "fmt"
)

func main() {
    var (
        environment string = "development"
    )

    flag.Func("environment", "Operating environment", func(flagValue string) error {
        for _, allowedValue := range []string{"development", "staging", "production"} {
            if flagValue == allowedValue {
                environment = flagValue
                return nil
            }
        }
        return errors.New(`must be one of "development", "staging" or "production"`)
    })

    flag.Parse()

    fmt.Printf("The operating environment is: %s\n", environment)
}

Making reusable helpers

If you find yourself repeating the same code in your flag.Func() functions, or the logic is getting too complex, it's possible to break it out into a reusable helper. For example, we could rewrite the example above to process our --environment flag via a generic enumFlag() function, like so:

package main

import (
    "flag"
    "fmt"
)

func main() {
    var (
        environment string = "development"
    )

    enumFlag(&environment, "environment", []string{"development", "staging", "production"}, "Operating environment")

    flag.Parse()

    fmt.Printf("The operating environment is: %s\n", environment)
}

func enumFlag(target *string, name string, safelist []string, usage string) {
    flag.Func(name, usage, func(flagValue string) error {
        for _, allowedValue := range safelist {
            if flagValue == allowedValue {
                *target = flagValue
                return nil
            }
        }

        return fmt.Errorf("must be one of %v", safelist)
    })
}