writer.go

package fake

import (
	"bytes"
	"fmt"
	"go/format"
	"io"
	"os"
	"sync"

	"github.com/rs/zerolog/log"
	"github.com/sonalys/fake/internal/files"
)

var pool = sync.Pool{
	New: func() any {
		newBuf := make([]byte, 0, 1024*10)
		return &newBuf
	},
}

func (g *Generator) GenerateFile(input string, interfaceNames ...string) []byte {
	parsedFile, err := g.ParseFile(input)
	if err != nil {
		log.Panic().Err(err).Msgf("failed to parse file: %s", input)
	}
	interfaces := parsedFile.ListInterfaces(interfaceNames...)
	if len(interfaces) == 0 {
		return nil
	}
	buf1 := pool.Get().(*[]byte)
	buf2 := pool.Get().(*[]byte)
	defer pool.Put(buf1)
	defer pool.Put(buf2)
	header := bytes.NewBuffer(*buf1)
	body := bytes.NewBuffer(*buf2)
	if g.MockPackageName == "" {
		g.MockPackageName = parsedFile.PkgName
	}
	writeHeader(header, g.MockPackageName)
	// Iterate through the declarations in the file
	for _, i := range interfaces {
		i.write(body)
	}
	// writeImports comes after interfaces because we only add external dependencies after generating interfaces.
	parsedFile.writeImports(header)
	header.ReadFrom(body)
	return formatCode(header.Bytes())
}

func writeHeader(w io.Writer, packageName string) {
	fmt.Fprintf(w, "// Code generated by fake. DO NOT EDIT.\n\n")
	fmt.Fprintf(w, "package %s\n\n", packageName)
}

func formatCode(in []byte) []byte {
	out, err := format.Source(in)
	if err != nil {
		log.Panic().Msgf("error formatting file: %v\n", err)
	}
	return out
}

func openOutputFile(input, output string) *os.File {
	outFile, err := files.CreateFileAndFolders(files.GenerateOutputFileName(input, output))
	if err != nil {
		log.Panic().Msgf("error creating mock file: %v\n", err)
	}
	return outFile
}