entc: introduce the entc.Extension api

examples/entcpkg/ent/entc.go

@@ -18,14 +18,15 @@ import (
 )
 
 func main() {
-	// A usage for custom templates with external functions.
-	// One template is defined in the option below, and the
-	// rest are provided with the `Templates` option.
-	opts := []entc.Option{
-		entc.TemplateFiles("template/stringer.tmpl"),
-		entc.Annotations(Annotation{StructTag: "rql"}),
+	ex, err := NewExtension("json")
+	if err != nil {
+		log.Fatalf("creating extension: %v", err)
 	}
-	err := entc.Generate("./schema", &gen.Config{
+	// A usage for custom options to configure the
+	opts := []entc.Option{
+		entc.Extensions(ex),
+	}
+	err = entc.Generate("./schema", &gen.Config{
 		Header: `
 			// Copyright 2019-present Facebook Inc. All rights reserved.
 			// This source code is licensed under the Apache 2.0 license found
@@ -34,20 +35,83 @@ func main() {
 			// Code generated by entc, DO NOT EDIT.
 		`,
 		Templates: []*gen.Template{
-			gen.MustParse(gen.NewTemplate("static").
-				Funcs(template.FuncMap{"title": strings.ToTitle}).
-				ParseFiles("template/static.tmpl")),
-			gen.MustParse(gen.NewTemplate("debug").
-				Funcs(template.FuncMap{"byName": byName}).
-				ParseFiles("template/debug.tmpl")),
+			// Custom templates can be provided by entc.Extension (see below),
+			// or by setting templates on the gen.Config object.
+			//
+			//	gen.MustParse(gen.NewTemplate("static").
+			//		Funcs(template.FuncMap{"title": strings.ToTitle}).
+			//		ParseFiles("template/static.tmpl")),
+			//
+		},
+		Hooks: []gen.Hook{
+			// Hooks can be provided by entc.Extension (see below),
+			// or by setting hooks on the gen.Config object.
+			//
+			// 	CustomHook1("config 1"),
+			// 	CustomHook2("config 2"),
+			//
 		},
-		Hooks: []gen.Hook{TagFields("json")},
 	}, opts...)
 	if err != nil {
 		log.Fatalf("running ent codegen: %v", err)
 	}
 }
 
+// Extension is an example implementation of entc.Extension.
+type Extension struct {
+	entc.DefaultExtension
+	tag       string
+	templates []*gen.Template
+}
+
+// NewExtension creates a new entc.Extension.
+func NewExtension(tag string) (*Extension, error) {
+	ex := &Extension{tag: tag}
+	t, err := gen.NewTemplate("static").
+		Funcs(template.FuncMap{"title": strings.ToTitle}).
+		ParseFiles("template/static.tmpl")
+	if err != nil {
+		return nil, err
+	}
+	ex.templates = append(ex.templates, t)
+	t, err = gen.NewTemplate("debug").
+		Funcs(template.FuncMap{"byName": byName}).
+		ParseFiles("template/debug.tmpl")
+	if err != nil {
+		return nil, err
+	}
+	ex.templates = append(ex.templates, t)
+	return ex, nil
+}
+
+// Templates of the extension.
+func (e *Extension) Templates() []*gen.Template {
+	return e.templates
+}
+
+// Hooks of the extension.
+func (e *Extension) Hooks() []gen.Hook {
+	return []gen.Hook{
+		TagFields(e.tag),
+	}
+}
+
+// Annotations of the extension.
+func (e *Extension) Annotations() []entc.Annotation {
+	return []entc.Annotation{
+		Annotation{StructTag: "rql"},
+	}
+}
+
+// Options provides additional options for the extension.
+func (e *Extension) Options() []entc.Option {
+	return []entc.Option{
+		entc.TemplateFiles("template/stringer.tmpl"),
+	}
+}
+
+var _ entc.Extension = (*Extension)(nil)
+
 // byName returns a node in the graph by its label/name.
 func byName(g *gen.Graph, name string) (*gen.Type, error) {
 	for _, n := range g.Nodes {