Skip to content

Commit b2dbc5d

Browse files
committed
Avoid reparsing templates and recompiling schemas
1 parent f738601 commit b2dbc5d

File tree

2 files changed

+81
-34
lines changed

2 files changed

+81
-34
lines changed

internal/remote_template.go

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010

1111
"github.com/chigopher/pathlib"
1212
"github.com/rs/zerolog"
13+
"github.com/vektra/mockery/v3/template"
1314
"github.com/xeipuuv/gojsonschema"
1415
)
1516

@@ -79,10 +80,10 @@ func (c *RemoteTemplateCache) Get(templateURL, schemaURL string) *RemoteTemplate
7980
}
8081

8182
type RemoteTemplate struct {
82-
templateURL string
83-
templateString string
84-
templateOnce sync.Once
85-
templateErr error
83+
templateURL string
84+
template template.Template
85+
templateOnce sync.Once
86+
templateErr error
8687

8788
schemaURL string
8889
schema *gojsonschema.Schema
@@ -99,21 +100,25 @@ func NewRemoteTemplate(templateURL, schemaURL string) *RemoteTemplate {
99100
}
100101
}
101102

102-
// Template will return the template string. It downloads the remote template once
103+
// Template returns the parsed template. It downloads the remote template once
103104
// and caches the result for future calls.
104-
func (r *RemoteTemplate) Template(ctx context.Context) (string, error) {
105+
func (r *RemoteTemplate) Template(ctx context.Context) (template.Template, error) {
105106
r.templateOnce.Do(func() {
106-
var err error
107-
r.templateString, err = download(ctx, r.templateURL)
107+
templateString, err := download(ctx, r.templateURL)
108108
if err != nil {
109109
r.templateErr = fmt.Errorf("downloading template: %w", err)
110+
return
111+
}
112+
r.template, err = template.New(templateString, r.templateURL)
113+
if err != nil {
114+
r.templateErr = fmt.Errorf("creating new template: %w", err)
110115
}
111116
})
112117

113-
return r.templateString, r.templateErr
118+
return r.template, r.templateErr
114119
}
115120

116-
// Schema returns the JSON Schema as a string. It downloads the remote schema once
121+
// Schema returns the compiled JSON Schema. It downloads the remote schema once
117122
// and caches the result for future calls.
118123
func (r *RemoteTemplate) Schema(ctx context.Context) (*gojsonschema.Schema, error) {
119124
log := zerolog.Ctx(ctx)

internal/template_generator.go

Lines changed: 66 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"go/types"
1313
"os"
1414
"strings"
15+
"sync"
1516

1617
"github.com/chigopher/pathlib"
1718
"github.com/rs/zerolog"
@@ -48,16 +49,58 @@ var (
4849

4950
var errBadHTTPStatus = errors.New("failed to download file")
5051

51-
var styleTemplates = map[string]string{
52-
"matryer": templateMatryer,
53-
"testify": templateTestify,
54-
"gomock": templateGomock,
52+
var styleTemplates = map[string]*StyleTemplate{
53+
"matryer": NewStyleTemplate("matryer", templateMatryer, templateMatryerJSONSchema),
54+
"testify": NewStyleTemplate("testify", templateTestify, templateTestifyJSONSchema),
55+
"gomock": NewStyleTemplate("gomock", templateGomock, templateGomockJSONSchema),
5556
}
5657

57-
var jsonSchemas = map[string]string{
58-
"matryer": templateMatryerJSONSchema,
59-
"testify": templateTestifyJSONSchema,
60-
"gomock": templateGomockJSONSchema,
58+
type StyleTemplate struct {
59+
name string
60+
61+
templateString string
62+
template template.Template
63+
templateOnce sync.Once
64+
templateErr error
65+
66+
schemaString string
67+
schema *gojsonschema.Schema
68+
schemaOnce sync.Once
69+
schemaErr error
70+
}
71+
72+
func NewStyleTemplate(name, templateString, schemaString string) *StyleTemplate {
73+
return &StyleTemplate{
74+
name: name,
75+
templateString: templateString,
76+
schemaString: schemaString,
77+
}
78+
}
79+
80+
// Template returns the parsed template.
81+
func (r *StyleTemplate) Template() (template.Template, error) {
82+
r.templateOnce.Do(func() {
83+
var err error
84+
r.template, err = template.New(r.templateString, r.name)
85+
if err != nil {
86+
r.templateErr = fmt.Errorf("creating new template: %w", err)
87+
}
88+
})
89+
90+
return r.template, r.templateErr
91+
}
92+
93+
// Schema returns the compiled JSON Schema.
94+
func (r *StyleTemplate) Schema() (*gojsonschema.Schema, error) {
95+
r.schemaOnce.Do(func() {
96+
var err error
97+
r.schema, err = gojsonschema.NewSchema(gojsonschema.NewStringLoader(r.schemaString))
98+
if err != nil {
99+
r.schemaErr = fmt.Errorf("creating JSON schema: %w", err)
100+
}
101+
})
102+
103+
return r.schema, r.schemaErr
61104
}
62105

63106
// findPkgPath returns the fully-qualified go import path of a given dir. The
@@ -333,7 +376,7 @@ func (g *TemplateGenerator) typeParams(ctx context.Context, tparams *types.TypeP
333376
}
334377

335378
// getTemplate returns the requested template and associated schema (if available).
336-
func (g *TemplateGenerator) getTemplate(ctx context.Context) (string, *gojsonschema.Schema, error) {
379+
func (g *TemplateGenerator) getTemplate(ctx context.Context) (template.Template, *gojsonschema.Schema, error) {
337380
log := zerolog.Ctx(ctx).With().Str("template", g.templateName).Str("schema", g.templateSchema).Logger()
338381
ctx = log.WithContext(ctx)
339382

@@ -345,33 +388,37 @@ func (g *TemplateGenerator) getTemplate(ctx context.Context) (string, *gojsonsch
345388
continue
346389
}
347390
remoteTemplate := g.remoteTemplateCache.Get(g.templateName, g.templateSchema)
348-
templateString, err := remoteTemplate.Template(ctx)
391+
templ, err := remoteTemplate.Template(ctx)
349392
if err != nil {
350393
log.Error().Msg("could not download template")
351-
return "", nil, fmt.Errorf("downloading template: %w", err)
394+
return template.Template{}, nil, fmt.Errorf("downloading template: %w", err)
352395
}
353396
if g.requireSchemaExists {
354397
schema, err = remoteTemplate.Schema(ctx)
355398
if err != nil {
356399
log.Error().Msg("could not get JSON schema")
357-
return "", nil, fmt.Errorf("downloading schema: %w", err)
400+
return template.Template{}, nil, fmt.Errorf("downloading schema: %w", err)
358401
}
359402
}
360403

361-
return templateString, schema, nil
404+
return templ, schema, nil
362405
}
363406

364407
// Embedded templates
365408
var styleExists bool
366-
templateString, styleExists := styleTemplates[g.templateName]
409+
styleTemplate, styleExists := styleTemplates[g.templateName]
367410
if !styleExists {
368-
return "", nil, fmt.Errorf("template '%s' does not exist", g.templateName)
411+
return template.Template{}, nil, fmt.Errorf("template '%s' does not exist", g.templateName)
412+
}
413+
templ, err := styleTemplate.Template()
414+
if err != nil {
415+
return template.Template{}, nil, err
369416
}
370-
schema, err := gojsonschema.NewSchema(gojsonschema.NewStringLoader(jsonSchemas[g.templateName]))
417+
schema, err := styleTemplate.Schema()
371418
if err != nil {
372-
return "", nil, fmt.Errorf("generating schema: %w", err)
419+
return template.Template{}, nil, err
373420
}
374-
return templateString, schema, nil
421+
return templ, schema, nil
375422
}
376423

377424
func validateSchema(ctx context.Context, data template.Data, schema *gojsonschema.Schema) error {
@@ -458,7 +505,7 @@ func (g *TemplateGenerator) Generate(
458505
data.SrcPkgQualifier = g.registry.SrcPkgName() + "."
459506
}
460507

461-
templateString, schema, err := g.getTemplate(ctx)
508+
templ, schema, err := g.getTemplate(ctx)
462509
if err != nil {
463510
log.Error().Msg("could not get template")
464511
return nil, fmt.Errorf("getting template: %w", err)
@@ -470,11 +517,6 @@ func (g *TemplateGenerator) Generate(
470517
}
471518
}
472519

473-
templ, err := template.New(templateString, g.templateName)
474-
if err != nil {
475-
return []byte{}, fmt.Errorf("creating new template: %w", err)
476-
}
477-
478520
var buf bytes.Buffer
479521
log.Debug().Msg("executing template")
480522
if err := templ.Execute(&buf, data); err != nil {

0 commit comments

Comments
 (0)