Using Go Generate

Posted on
go golang meta programming

Recently I needed to define some types who’s main purpose were to simplify the use of some exported package functions. The types represented hundreds of domain specific constants, with each type’s constants made available in separate text documents. This was not an impossible task but certainly not one I would enjoy doing. Fortunately, Go 1.4 introduced a new command, generate, that would make it easier to create these types.

I searched through the official Go sources to find an example and came across one in /x/net/html/atom. The usage of go generate was to create an efficient lookup table of integer codes to strings. My needs were slightly simpler, so I adapted the process to read from a file and omitted the optimizations.

From start to finish our pseudocode looks like this

for each T in custom types
	read constants from T.txt
	fill a buffer with T definition and associated constants
	format buffer with go format
	save buffer to T.go

This process occurs in a main package excluded from the build process but included with our source in the package for which it will generate code. To demonstrate a practical example, we’ll generate HTML5 and SVG tags from a text file.

Package main

The main package is defined below in a file called gen.go. Note the exclusion from the build process and the go:generate comment, indicating to the generate command what it should do. In our case, we want it to execute go run gen.go.

// +build ignore

//go:generate go run gen.go
package main

// imports 

func parseFile(path string) ([]string, error) {
	// omitted
}

func main() {
	// code..
}

In func main() we begin by parsing the text file into a slice of strings.

func main() {
	tags, err := parseFile("tags.txt")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	// ...
}

The next section covers the actual source code we will insert into our generated .go file. As of Go 1.9, we should start our machine generated code with a comment that matches the following regex: ^// Code generated .* DO NOT EDIT.$.

We’ll include the commands used to generate the code in our comment. We’ll also use the package name tags and define a type Tag that implements the Stringer interface.

	var buf bytes.Buffer
	fmt.Fprint(&buf, "// Code generated by go generate gen.go DO NOT EDIT.\n\n")
	fmt.Fprint(&buf, "//go:generate go run gen.go\n\n")
	fmt.Fprint(&buf, "package tags\n\n")
	fmt.Fprint(&buf, "type Tag string\n\n")
	fmt.Fprint(&buf, "func (t Tag) String() string { return string(t) }\n\n")

Our next step will be to generate a giant const block of tags, based on input from our text file.

	fmt.Fprint(&buf, "const (\n")
	for _, tag := range tags {
		fmt.Fprintf(&buf, "\t%s\tTag = \"%s\"\n", strings.Title(tag), tag)
	}
	fmt.Fprint(&buf, "\tUnknown\tTag = \"\"\n)")

Note the use of strings.Title to make sure our tag name is exported.

The last step in the process is to make sure our code is run through gofmt and then written to disk. The call to format.Source() will handle formatting or output any errors that occurred during formatting.

	b, err := format.Source(buf.Bytes())
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	if err := ioutil.WriteFile("tags.go", b, 0644); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

Running go generate gen.go will create a file called tags.go in the same directory. Check out the source code for this post on my Github.

Thanks for reading