oapi-codegen · deepmap-marcinr · Feb 19, 2020 · Feb 18, 2020 · Feb 19, 2020
diff --git a/README.md b/README.md
@@ -454,4 +454,16 @@ Go file.
 Afterwards you should run `go generate ./...`, and the templates will be updated
  accordingly.
 
-
+Alternatively, you can provide custom templates to override built-in ones using
+the `-templates` flag specifying a path to a directory containing templates
+files. These files **must** be named identically to built-in template files
+(see `pkg/codegen/templates/*.tmpl` in the source code), and will be interpreted
+on-the-fly at run time. Example:
+
+    $ ls -1 my-templates/
+    client.tmpl
+    typedef.tmpl
+    $ oapi-codegen \
+        -templates my-templates/ \
+        -generate types,client \
+        petstore-expanded.yaml
diff --git a/cmd/oapi-codegen/oapi-codegen.go b/cmd/oapi-codegen/oapi-codegen.go
@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"path"
 	"path/filepath"
 	"strings"
 
@@ -32,18 +33,20 @@ func errExit(format string, args ...interface{}) {
 
 func main() {
 	var (
-		packageName string
-		generate    string
-		outputFile  string
-		includeTags string
-		excludeTags string
+		packageName  string
+		generate     string
+		outputFile   string
+		includeTags  string
+		excludeTags  string
+		templatesDir string
 	)
 	flag.StringVar(&packageName, "package", "", "The package name for generated code")
 	flag.StringVar(&generate, "generate", "types,client,server,spec",
 		`Comma-separated list of code to generate; valid options: "types", "client", "chi-server", "server", "skip-fmt", "spec"`)
 	flag.StringVar(&outputFile, "o", "", "Where to output generated code, stdout is default")
 	flag.StringVar(&includeTags, "include-tags", "", "Only include operations with the given tags. Comma-separated list of tags.")
 	flag.StringVar(&excludeTags, "exclude-tags", "", "Exclude operations that are tagged with the given tags. Comma-separated list of tags.")
+	flag.StringVar(&templatesDir, "templates", "", "Path to directory containing user templates")
 	flag.Parse()
 
 	if flag.NArg() < 1 {
@@ -95,6 +98,12 @@ func main() {
 		errExit("error loading swagger spec\n: %s", err)
 	}
 
+	templates, err := loadTemplateOverrides(templatesDir)
+	if err != nil {
+		errExit("error loading template overrides: %s\n", err)
+	}
+	opts.UserTemplates = templates
+
 	code, err := codegen.Generate(swagger, packageName, opts)
 	if err != nil {
 		errExit("error generating code: %s\n", err)
@@ -125,3 +134,26 @@ func splitCSVArg(input string) []string {
 	}
 	return args
 }
+
+func loadTemplateOverrides(templatesDir string) (map[string]string, error) {
+	var templates = make(map[string]string)
+
+	if templatesDir == "" {
+		return templates, nil
+	}
+
+	files, err := ioutil.ReadDir(templatesDir)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, f := range files {
+		data, err := ioutil.ReadFile(path.Join(templatesDir, f.Name()))
+		if err != nil {
+			return nil, err
+		}
+		templates[f.Name()] = string(data)
+	}
+
+	return templates, nil
+}
diff --git a/pkg/codegen/codegen.go b/pkg/codegen/codegen.go
@@ -32,14 +32,15 @@ import (
 
 // Options defines the optional code to generate.
 type Options struct {
-	GenerateChiServer  bool     // GenerateChiServer specifies whether to generate chi server boilerplate
-	GenerateEchoServer bool     // GenerateEchoServer specifies whether to generate echo server boilerplate
-	GenerateClient     bool     // GenerateClient specifies whether to generate client boilerplate
-	GenerateTypes      bool     // GenerateTypes specifies whether to generate type definitions
-	EmbedSpec          bool     // Whether to embed the swagger spec in the generated code
-	SkipFmt            bool     // Whether to skip go fmt on the generated code
-	IncludeTags        []string // Only include operations that have one of these tags. Ignored when empty.
-	ExcludeTags        []string // Exclude operations that have one of these tags. Ignored when empty.
+	GenerateChiServer  bool              // GenerateChiServer specifies whether to generate chi server boilerplate
+	GenerateEchoServer bool              // GenerateEchoServer specifies whether to generate echo server boilerplate
+	GenerateClient     bool              // GenerateClient specifies whether to generate client boilerplate
+	GenerateTypes      bool              // GenerateTypes specifies whether to generate type definitions
+	EmbedSpec          bool              // Whether to embed the swagger spec in the generated code
+	SkipFmt            bool              // Whether to skip go fmt on the generated code
+	IncludeTags        []string          // Only include operations that have one of these tags. Ignored when empty.
+	ExcludeTags        []string          // Exclude operations that have one of these tags. Ignored when empty.
+	UserTemplates      map[string]string // Override built-in templates from user-provided files
 }
 
 type goImport struct {
@@ -99,6 +100,16 @@ func Generate(swagger *openapi3.Swagger, packageName string, opts Options) (stri
 		return "", errors.Wrap(err, "error parsing oapi-codegen templates")
 	}
 
+	// Override built-in templates with user-provided versions
+	for _, tpl := range t.Templates() {
+		if _, ok := opts.UserTemplates[tpl.Name()]; ok {
+			utpl := t.New(tpl.Name())
+			if _, err := utpl.Parse(opts.UserTemplates[tpl.Name()]); err != nil {
+				return "", errors.Wrapf(err, "error parsing user-provided template %q", tpl.Name())
+			}
+		}
+	}
+
 	ops, err := OperationDefinitions(swagger)
 	if err != nil {
 		return "", errors.Wrap(err, "error creating operation definitions")

diff --git a/pkg/codegen/codegen_test.go b/pkg/codegen/codegen_test.go
@@ -59,6 +59,37 @@ func TestExamplePetStoreCodeGeneration(t *testing.T) {
 	assert.Len(t, problems, 0)
 }
 
+func TestExamplePetStoreCodeGenerationWithUserTemplates(t *testing.T) {
+
+	userTemplates := map[string]string{"typedef.tmpl": "//blah"}
+
+	// Input vars for code generation:
+	packageName := "api"
+	opts := Options{
+		GenerateTypes: true,
+		UserTemplates: userTemplates,
+	}
+
+	// Get a spec from the example PetStore definition:
+	swagger, err := examplePetstore.GetSwagger()
+	assert.NoError(t, err)
+
+	// Run our code generation:
+	code, err := Generate(swagger, packageName, opts)
+	assert.NoError(t, err)
+	assert.NotEmpty(t, code)
+
+	// Check that we have valid (formattable) code:
+	_, err = format.Source([]byte(code))
+	assert.NoError(t, err)
+
+	// Check that we have a package:
+	assert.Contains(t, code, "package api")
+
+	// Check that the built-in template has been overriden
+	assert.Contains(t, code, "//blah")
+}
+
 func TestExamplePetStoreParseFunction(t *testing.T) {
 
 	bodyBytes := []byte(`{"id": 5, "name": "testpet", "tag": "cat"}`)