Skip to content

Commit

Permalink
Merge pull request #24 from raskyld/13-handle-custom-headers
Browse files Browse the repository at this point in the history 
Custom headers
  • Loading branch information
@raskyld
raskyld authored Oct 24, 2023
2 parents 5c3bd3b + 5bbbf46 commit c5f1ae1
Show file tree
Hide file tree
Showing 4 changed files with 83 additions and 30 deletions.
11 changes: 9 additions & 2 deletions cmd/gen.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ import (
"github.com/spf13/cobra"
"path/filepath"
"sigs.k8s.io/controller-tools/pkg/genall"
"strconv"
"time"
)

func NewGenerate(ctx *Context) *cobra.Command {
Expand All @@ -40,6 +42,9 @@ func NewGenerate(ctx *Context) *cobra.Command {
}

func NewGenerateGo(ctx *Context) *cobra.Command {
var headerFile string
var year string

genFuncGo := &cobra.Command{
Use: "go [internalPkgPath] [internalPkgName]",
Short: "Generate the Go code to integrate with Tekton",
Expand Down Expand Up @@ -67,14 +72,14 @@ The code generated for your main package will be written in zz_generated.tektask

var genFunc, genInternal genall.Generator

genFuncPtr, err := gengo.NewGoFunc(ctx.Logger)
genFuncPtr, err := gengo.NewGoFunc(ctx.Logger, headerFile, year)
if err != nil {
return err
}

genFunc = genFuncPtr

genInternalPtr, err := gengo.NewGoInternal(ctx.Logger, outputPkgName)
genInternalPtr, err := gengo.NewGoInternal(ctx.Logger, outputPkgName, headerFile, year)
if err != nil {
return err
}
Expand Down Expand Up @@ -107,6 +112,8 @@ The code generated for your main package will be written in zz_generated.tektask
},
}

genFuncGo.Flags().StringVarP(&headerFile, "headerfile", "b", "", "Path to a boilerplate header file to put at the top of any generated go code. TIPS: ' YEAR' will be replaced with the current year!")
genFuncGo.Flags().StringVar(&year, "year", strconv.Itoa(time.Now().Year()), "Which year should be written in the headerfile")
return genFuncGo
}

Expand Down
73 changes: 52 additions & 21 deletions internal/gengo/gen_func.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,45 +19,61 @@ package gengo
import (
"bytes"
"fmt"
ttmarkers "github.com/Raskyld/go-tektasker/pkg/markers"
"go/ast"
"log/slog"
"strings"
"text/template"

ttmarkers "github.com/Raskyld/go-tektasker/pkg/markers"
"sigs.k8s.io/controller-tools/pkg/genall"
"sigs.k8s.io/controller-tools/pkg/markers"
"text/template"
)

const FuncName = "func"
const FuncTpl = `{{template "%s" .Header}}
const FuncTpl = `{{template "%s" .GoHeaderArgs}}
{{- range $tplName, $args := .Templates}}
{{- range $tplName, $args := .TemplatesArgs}}
{{- range $args}}
{{CallTemplate $tplName .}}
{{- end}}
{{- end}}
`

// TaskGoFuncGenerator is the generator in charge of generating Go files for tasks
// ATM, it generates both the internal Tektasker packages and makes user types implement
// Interface awaited by Helper functions
type TaskGoFuncGenerator struct {
Logger *slog.Logger

Template *template.Template

// HeaderFile path to add at the top of every generated go file
HeaderFile string

// Year to use for header copyright notice
Year string
}

// PerTemplateArgs maps each registered template to a mapping of param or result name to template args
// templateName -> Param/Result Name -> Param/ResultArgs
// We then will execute each template for each registered param or name with
// the associated args
type PerTemplateArgs map[string]map[string]interface{}

type FuncArgs struct {
Header GoHeaderArgs
Templates PerTemplateArgs
// GoHeaderArgs is the list of arguments we will pass to the GoHeader template invoked on every generated
// go file
GoHeaderArgs GoHeaderArgs

// TemplatesArgs is a list of args to pass to the registered templates that is indexed by template,
// then by param or result name
TemplatesArgs PerTemplateArgs
}

func NewGoFunc(logger *slog.Logger) (*TaskGoFuncGenerator, error) {
func NewGoFunc(logger *slog.Logger, headerFile, year string) (*TaskGoFuncGenerator, error) {
g := &TaskGoFuncGenerator{
Logger: logger.With("generator", "goFunc"),
Template: &template.Template{},
Logger: logger.With("generator", "goFunc"),
Template: &template.Template{},
HeaderFile: headerFile,
Year: year,
}

// NB(raskyld): this hack is needed to call template
Expand Down Expand Up @@ -86,6 +102,7 @@ func NewGoFunc(logger *slog.Logger) (*TaskGoFuncGenerator, error) {
return g, nil
}

// RegisterTemplate with a unique name which can then latter be used by the generator
func (g *TaskGoFuncGenerator) RegisterTemplate(name string, tpl string) *TaskGoFuncGenerator {
nt, err := g.Template.New(name).Parse(tpl)
if err != nil {
Expand All @@ -101,6 +118,17 @@ func (*TaskGoFuncGenerator) RegisterMarkers(into *markers.Registry) error {
}

func (g *TaskGoFuncGenerator) Generate(ctx *genall.GenerationContext) error {
var headerText string

if g.HeaderFile != "" {
buf, err := ctx.ReadFile(g.HeaderFile)
if err != nil {
return err
}

headerText = strings.ReplaceAll(string(buf), " YEAR", " "+g.Year)
}

for _, pkg := range ctx.Roots {
logger := g.Logger.With("pkg", pkg.Name)
logger.Debug("starting collecting")
Expand All @@ -118,6 +146,7 @@ func (g *TaskGoFuncGenerator) Generate(ctx *genall.GenerationContext) error {
continue
}

// Prepare the per-template per-param/result args mappings
perTemplateArgs := PerTemplateArgs(make(map[string]map[string]interface{}))
for _, t := range g.Template.Templates() {
// NB(raskyld): we need to exclude the main template as
Expand Down Expand Up @@ -151,18 +180,17 @@ func (g *TaskGoFuncGenerator) Generate(ctx *genall.GenerationContext) error {
// TODO(raskyld):
// add a way to skip creating the Unmarshal method so our
// users can create a custom unmarshaler for their complex types

mapToAdd := ParamFuncUnmarshalJSONName
funcTemplateToUse := ParamFuncUnmarshalJSONName

// NB(raskyld): this is for ease of use when we have a type made of string
// otherwise, we just expect the value to be valid JSON
if ident, ok := info.RawSpec.Type.(*ast.Ident); ok {
if ident.Name == "string" {
mapToAdd = ParamFuncUnmarshalSimpleName
funcTemplateToUse = ParamFuncUnmarshalSimpleName
}
}

perTemplateArgs[mapToAdd][param.Name] = ParamFuncArgs{
perTemplateArgs[funcTemplateToUse][param.Name] = ParamFuncArgs{
ParamName: param.Name,
ParamType: info.Name,
}
Expand All @@ -186,23 +214,26 @@ func (g *TaskGoFuncGenerator) Generate(ctx *genall.GenerationContext) error {
ResultType: info.Name,
}

mapToAdd := ResultFuncMarshalJSONName
funcTemplateToUse := ResultFuncMarshalJSONName

// NB(raskyld): this is for ease of use when we have a type made of string
// otherwise, we just expect the value to be valid JSON
if ident, ok := info.RawSpec.Type.(*ast.Ident); ok {
if ident.Name == "string" {
mapToAdd = ResultFuncMarshalSimpleName
funcTemplateToUse = ResultFuncMarshalSimpleName
}
}

perTemplateArgs[mapToAdd][result.Name] = ResultFuncArgs{
perTemplateArgs[funcTemplateToUse][result.Name] = ResultFuncArgs{
ResultName: result.Name,
ResultType: info.Name,
}
}
}
})
if err != nil {
return err
}

output, err := ctx.OutputRule.Open(pkg, "zz_generated.tektasker.go")
if err != nil {
Expand All @@ -215,12 +246,12 @@ func (g *TaskGoFuncGenerator) Generate(ctx *genall.GenerationContext) error {
}

err = g.Template.ExecuteTemplate(output, FuncName, FuncArgs{
Header: GoHeaderArgs{
GoHeaderArgs: GoHeaderArgs{
PkgName: pkg.Name,
Header: "// generated code",
Header: headerText,
ImportPaths: importPaths,
},
Templates: perTemplateArgs,
TemplatesArgs: perTemplateArgs,
})

if err != nil {
Expand Down
25 changes: 19 additions & 6 deletions internal/gengo/gen_internal.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,22 +22,25 @@ import (
"log/slog"
"sigs.k8s.io/controller-tools/pkg/genall"
"sigs.k8s.io/controller-tools/pkg/markers"
"strings"
"text/template"
)

type TaskGoInternalGenerator struct {
Logger *slog.Logger

Template *template.Template

Logger *slog.Logger
Template *template.Template
PackageName string
HeaderFile string
Year string
}

func NewGoInternal(logger *slog.Logger, pkgName string) (*TaskGoInternalGenerator, error) {
func NewGoInternal(logger *slog.Logger, pkgName, headerFile, year string) (*TaskGoInternalGenerator, error) {
g := &TaskGoInternalGenerator{
Logger: logger.With("generator", "goInternal"),
Template: &template.Template{},
PackageName: pkgName,
HeaderFile: headerFile,
Year: year,
}

g.RegisterTemplate(GoHeaderName, GoHeaderTpl).
Expand All @@ -63,10 +66,20 @@ func (*TaskGoInternalGenerator) RegisterMarkers(into *markers.Registry) error {

func (g *TaskGoInternalGenerator) Generate(ctx *genall.GenerationContext) error {
var headerBytes bytes.Buffer
var headerText string

if g.HeaderFile != "" {
buf, err := ctx.ReadFile(g.HeaderFile)
if err != nil {
return err
}

headerText = strings.ReplaceAll(string(buf), " YEAR", " "+g.Year)
}

headerArgs := GoHeaderArgs{
PkgName: g.PackageName,
Header: "// generated code",
Header: headerText,
ImportPaths: []string{
"errors",
"fmt",
Expand Down
4 changes: 3 additions & 1 deletion internal/gengo/go_header.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,9 @@ package gengo

const GoHeaderName = "go.header"

const GoHeaderTpl = `{{.Header}}
const GoHeaderTpl = `{{- with .Header -}}
{{.}}
{{end -}}
package {{.PkgName}}
Expand Down

0 comments on commit c5f1ae1

Please sign in to comment.