A taste of Go code generator magic: a quick guide to getting started

package db

type DB struct {
    // ... some private fields
}

func (d *DB) InsertDocuments(documents []*DocumentParams) ([]*Document, error) {
    // ...
}

func (d *DB) DeleteDocument(id string) error {
    // ...
}

var globalDB *DB

func SetGlobalDB(value *DB) {
    globalDB = value
}

// Boilerplate code: candidates for code generation

func InsertDocuments(documents []*DocumentParams) ([]*Document, error) {
    if globalDB == nil {
        panic("globalDB is not set")
    }

    return globalDB.InsertDocuments(documents)
}

func DeleteDocument(id string) error {
    if globalDB == nil {
        panic("globalDB is not set")
    }

    return globalDB.DeleteDocument(id)
}

package main

import (
	"log"

	"github.com/dave/jennifer/jen"
)

func main() {
	f := jen.NewFile("main")

    // A comment to mark the file as generated
	f.PackageComment("Code generated by generator, DO NOT EDIT.")

    f.Func().Id("Hello").Block(
        jen.Qual("fmt", "Println").Call(jen.Lit("Hello from generated code!")),
    )

	targetFilename := "main_gen.go"

	if err := f.Save(targetFilename); err != nil {
		log.Fatal(err)
	}
}

// Code generated by generator, DO NOT EDIT.
package main

import "fmt"

func Hello() {
        fmt.Println("Hello from generated code!")
}

package main

import (
	"go/types"
    "log"
    "strings"

	"golang.org/x/tools/go/packages"
	"github.com/dave/jennifer/jen"
)

func main() {
    // 1. Load the package metadata and get all public methods of a given type.

    if len(os.Args) < 3 {
        log.Fatal("Usage: //go:generate go run gen_functions.go <full package name> <type name>")
    }

    pkgName := os.Args[1]
    typeName := os.Args[2]

    // Load package metadata. We need it for accessing typeName's methods.
    pkg := loadPackage(pkgName)

	typeObject := pkg.Types.Scope().Lookup(typeName)
	if typeObject == nil {
		log.Fatalf("unable to find type %s", typeName)
	}

    // We will support only pointer semantics
	typeObjectMeta := types.NewPointer(typeObject.Type())

	methods := types.NewMethodSet(typeObjectMeta)
	for i := 0; i < methods.Len(); i++ {
		method := methods.At(i).Obj()

		// Skip private methods
		if !method.Exported() {
			continue
		}

        // 2. Generate the functions.
        // ...
    }

    // 3. Put the generated code
    // ...

    // 4. Save the results to a file
    // ...
}

func loadPackage(path string) *packages.Package {
	cfg := &packages.Config{Mode: packages.NeedTypes}
	pkgs, err := packages.Load(cfg, path)
	if err != nil {
		log.Fatalf("failed to load package '%s' for inspection: %v", path, err)
	}
	if packages.PrintErrors(pkgs) > 0 {
		log.Fatalf("package errors: %v", packages.PrintErrors(pkgs))
	}

	return pkgs[0]
}

package main

// ...

func main() {
    // 1. Load the package metadata and get all public methods of a given type.
    // ...

    varName := "global" + typeName
    imports := make(map[string]struct{})
    funcs := jen.Empty()

	for i := 0; i < methods.Len(); i++ {
        // ...

        // 2. Generate the functions.
		signature := method.Type().(*types.Signature)
		signatureParams := signature.Params()

		params := make([]jen.Code, 0, signatureParams.Len())
		paramNames := make([]jen.Code, 0, signatureParams.Len())

		for j := 0; j < signatureParams.Len(); j++ {
			param := signatureParams.At(j)
			paramType := param.Type().String()

			// Remove current package prefix
			paramType = strings.ReplaceAll(paramType, packageName+".", "")

			// Remove prefixes of imported packages
			paramType = extractImports(paramType, imports)

			params = append(params, jen.Id(param.Name()).Id(paramType))
			paramNames = append(paramNames, jen.Id(param.Name()))
		}

		results := signature.Results()
		returnTypes := make([]string, 0, results.Len())

		for j := 0; j < results.Len(); j++ {
			resultType := results.At(j).Type().String()

			// Remove current package prefix
			resultType = strings.ReplaceAll(resultType, packageName+".", "")

			// Remove prefixes of imported packages
			resultType = extractImports(resultType, imports)

			returnTypes = append(returnTypes, resultType)
		}

        var returnType string
		var lastStatement jen.Code = jen.Id(varName).Dot(method.Name()).Call(paramNames...)

		if len(returnTypes) > 0 {
			returnType = "(" + strings.Join(returnTypes, ", ") + ")"
			lastStatement = jen.Return(lastStatement)
		}

        // Generate the proxy function
		funcs.Comment(method.Name() + " generated.")
		funcs.Line()
		funcs.Func().Id(method.Name()).Params(params...).Id(returnType).Block(
            jen.If(jen.Id(varName).Op("==").Id("nil")).Block(
				jen.Panic(jen.Lit(varName + "  is not set.")),
			),
			jen.Line(),
			lastStatement,
		)
		funcs.Line()
    }

    // 3. Put the generated code
    // ...

    // 4. Save the results to a file
    // ...
}


var rePkgReferencedType = regexp.MustCompile(`([^*\[\]]+)\.[^.]+`)

// Remove imported packages prefixes and accumulate missing imports.
//
//	  github.com/somelib/name.Type   -> name.Type
//    []github.com/somelib/name.Item -> []name.Item
func extractImports(paramType string, imports map[string]struct{}) string {
	if strings.Index(paramType, ".") < 0 {
		return paramType
	}

	submatches := rePkgReferencedType.FindStringSubmatch(paramType)

	imports[submatches[1]] = struct{}{}

	for packageName := range imports {
		paramType = strings.ReplaceAll(paramType, filepath.Dir(packageName)+"/", "")
	}

	return paramType
}

package main

// ...

func main() {
    // 1. Load the package metadata and get all public methods of a given type.
    // ...

    // 2. Generate the functions.
    // ...

    // 3. Put the generated code

    // NOTE: `jen` has a Qual function that automatically adds new imports
    // but here we parse imports ourselves so we have to add them manually.
	if len(imports) > 0 {
		importsList := jen.Empty()
		for importPackage := range imports {
			importsList.Add(
				jen.Lit(importPackage),
				jen.Line(),
			)
		}

		f.Comment("Import missing dependencies")
		f.Add(jen.Id("import").Parens(importsList))
	}

    // Declare a global variable for our singleton
	f.Id("var").Id(varName).Id(typeName)

    // Add a function to set the global variable
	f.Func().Id("SetGlobal"+typeName).Params(jen.Id("value").Id(typeName)).Block(
		jen.Id(varName).Op("=").Id("value"),
	)

    // Add generated proxy functions
	f.Add(funcs)

    // 4. Save the results to a file

	filename := os.Getenv("GOFILE")
	ext := filepath.Ext(filename)
	filename, _ = strings.CutSuffix(filename, ext)
	filename = filename + "_gen.go"

	if err := f.Save(filename); err != nil {
		log.Fatal(err)
	}
}