main_file.go

package codegen

import (
	"os"
	"os/exec"
	"strings"
)

type File struct {
	imports *importsBlock
	stmts   []Block
}

// NewFile adds a package name and a comment that the file is auto-generated
func NewFile(packageName, generatorName string) *File {
	f := &File{stmts: make([]Block, 3), imports: newImportsBlock()}

	f.stmts[0] = newCommentF("Code generated by %s. DO NOT EDIT.", generatorName)
	f.stmts[1] = pkg(packageName)
	f.stmts[2] = f.imports

	return f
}

// Save creates a new file for the generated code
func (f *File) Save(filePath string) error {
	if out, err := os.Create(filePath); err != nil {
		return err
	} else {
		defer out.Close()

		if err = out.Truncate(0); err != nil {
			return err
		} else if _, err = out.Seek(0, 0); err != nil {
			return err
		} else if _, err = out.WriteString(f.GoString()); err != nil {
			return err
		}

		exec.Command("gofmt", "-w", filePath).Run()
		return nil
	}
}

// AddImport adds a new import statement to the import block
func (f *File) AddImport(path string) {
	f.imports.AddImport(path)
}

// AddImportAlias adds a new import statement with its package alias to the import block
func (f *File) AddImportAlias(path, alias string) {
	f.imports.AddImportAlias(path, alias)
}

func (f *File) GoString() string {
	var sb strings.Builder

	for _, stmt := range f.stmts {
		stmt.write(&sb)
	}

	return sb.String()
}

func (f *File) append(blk Block) {
	f.stmts = append(f.stmts, blk)
}