jirfag/go-queryset

package generator

import (
    "bytes"
    "context"
    "fmt"
    "io"
    "log"
    "os"
    "path/filepath"

    "github.com/jirfag/go-queryset/internal/parser"
    "github.com/pkg/errors"
    "golang.org/x/tools/imports"
)

type Generator struct {
    StructsParser *parser.Structs
}

// Generate generates output file with querysets
func (g Generator) Generate(ctx context.Context, inFilePath, outFilePath string) error {
    parsedFile, err := g.StructsParser.ParseFile(ctx, inFilePath)
    if err != nil {
        return errors.Wrapf(err, "can't parse file %s to get structs", inFilePath)
    }

    var r io.Reader
    r, err = GenerateQuerySetsForStructs(parsedFile.Types, parsedFile.Structs)
    if err != nil {
        return errors.Wrap(err, "can't generate query sets")
    }

    if r == nil {
        return fmt.Errorf("no structs to generate query set in %s", inFilePath)
    }

    if err = g.writeQuerySetsToOutput(r, parsedFile.PackageName, outFilePath); err != nil {
        return errors.Wrapf(err, "can't save query sets to out file %s", outFilePath)
    }

    var absOutPath string
    absOutPath, err = filepath.Abs(outFilePath)
    if err != nil {
        absOutPath = outFilePath
    }

    log.Printf("successfully wrote querysets to %s", absOutPath)
    return nil
}

func (g Generator) writeQuerySetsToOutput(r io.Reader, packageName, outFile string) error {
    const hdrTmpl = `%s
    package %s

import (
    "errors"
    "fmt"
    "strings"
    "time"

    "github.com/jinzhu/gorm"
)
`

    // https://golang.org/s/generatedcode
    const genHdr = `// Code generated by go-queryset. DO NOT EDIT.`

    var buf bytes.Buffer
    pkgName := fmt.Sprintf(hdrTmpl, genHdr, packageName)
    if _, err := buf.WriteString(pkgName); err != nil {
        return errors.Wrap(err, "can't write hdr string into buf")
    }
    if _, err := io.Copy(&buf, r); err != nil {
        return errors.Wrap(err, "can't write to buf")
    }

    formattedRes, err := imports.Process(outFile, buf.Bytes(), nil)
    if err != nil {
        if os.Getenv("GOQUERYSET_DEBUG_IMPORTS_ERRORS") == "1" {
            log.Printf("Can't format generated file: %s", err)
            formattedRes = buf.Bytes()
        } else {
            return errors.Wrap(err, "can't format generated file")
        }
    }

    var outF *os.File
    outF, err = os.OpenFile(outFile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0640) // nolint: gas
    if err != nil {
        return fmt.Errorf("can't open out file: %s", err)
    }
    defer func() {
        if e := outF.Close(); e != nil {
            log.Printf("can't close file: %s", e)
        }
    }()

    if _, err = outF.Write(formattedRes); err != nil {
        return errors.Wrap(err, "can't write to out file")
    }

    return nil
}