oapi-codegen · chanced · Aug 5, 2022 · Aug 6, 2022 · Aug 6, 2022
diff --git a/pkg/codegen/codegen.go b/pkg/codegen/codegen.go
@@ -20,6 +20,8 @@ import (
 	"embed"
 	"fmt"
 	"io/fs"
+	"io/ioutil"
+	"path"
 	"runtime/debug"
 	"sort"
 	"strings"
@@ -122,12 +124,30 @@ func Generate(spec *openapi3.T, opts Configuration) (string, error) {
 	if err != nil {
 		return "", fmt.Errorf("error parsing oapi-codegen templates: %w", err)
 	}
+	templateOverrides := make(map[string]string, len(opts.OutputOptions.UserTemplates))
+
+	// This loads user-provided templates by walking the supplied directory
+	if opts.OutputOptions.UserTemplatesDir != "" {
+		overrides, err := loadTemplateOverrides(opts.OutputOptions.UserTemplatesDir)
+		if err != nil {
+			return "", fmt.Errorf("error parsing user templates directory: %w", err)
+		}
+		for k, v := range overrides {
+			templateOverrides[k] = v
+		}
+	}
+
+	// This loads the user-templates, overriding any templates that were loaded
+	// from the user-templates-dir
+	for k, v := range opts.OutputOptions.UserTemplates {
+		templateOverrides[k] = v
+	}
 
 	// Override built-in templates with user-provided versions
 	for _, tpl := range t.Templates() {
-		if _, ok := opts.OutputOptions.UserTemplates[tpl.Name()]; ok {
+		if override, ok := templateOverrides[tpl.Name()]; ok {
 			utpl := t.New(tpl.Name())
-			if _, err := utpl.Parse(opts.OutputOptions.UserTemplates[tpl.Name()]); err != nil {
+			if _, err := utpl.Parse(override); err != nil {
 				return "", fmt.Errorf("error parsing user-provided template %q: %w", tpl.Name(), err)
 			}
 		}
@@ -896,3 +916,39 @@ func GetResponsesImports(responses map[string]*openapi3.ResponseRef) (map[string
 	}
 	return res, nil
 }
+
+func loadTemplateOverrides(templatesDir string) (map[string]string, error) {
+	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 {
+		// Recursively load subdirectory files, using the path relative to the templates
+		// directory as the key. This allows for overriding the files in the service-specific
+		// directories (e.g. echo, chi, etc.).
+		if f.IsDir() {
+			subFiles, err := loadTemplateOverrides(path.Join(templatesDir, f.Name()))
+			if err != nil {
+				return nil, err
+			}
+			for subDir, subFile := range subFiles {
+				templates[path.Join(f.Name(), subDir)] = subFile
+			}
+			continue
+		}
+		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_test.go b/pkg/codegen/codegen_test.go
@@ -28,7 +28,6 @@ func checkLint(t *testing.T, filename string, code []byte) {
 }
 
 func TestExamplePetStoreCodeGeneration(t *testing.T) {
-
 	// Input vars for code generation:
 	packageName := "api"
 	opts := Configuration{
@@ -73,7 +72,6 @@ func TestExamplePetStoreCodeGeneration(t *testing.T) {
 }
 
 func TestExamplePetStoreCodeGenerationWithUserTemplates(t *testing.T) {
-
 	userTemplates := map[string]string{"typedef.tmpl": "//blah"}
 
 	// Input vars for code generation:
@@ -108,8 +106,70 @@ func TestExamplePetStoreCodeGenerationWithUserTemplates(t *testing.T) {
 	assert.Contains(t, code, "//blah")
 }
 
-func TestExamplePetStoreParseFunction(t *testing.T) {
+func TestExamplePetStoreCodeGenerationWithUserTemplatesDir(t *testing.T) {
+	// Input vars for code generation:
+
+	opts := Configuration{
+		PackageName: "api",
+		Generate: GenerateOptions{
+			ChiServer: true,
+		},
+		OutputOptions: OutputOptions{
+			UserTemplatesDir: "test_templates",
+		},
+	}
+
+	// Get a spec from the example PetStore definition:
+	swagger, err := examplePetstore.GetSwagger()
+	assert.NoError(t, err)
+
+	// Run our code generation:
+	code, err := Generate(swagger, 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, "//__placeholder__")
+
+	// Tests that `user-templates` takes precedence over
+	// `user-templates-dir`
+
+	// Input vars for code generation:
+	opts = Configuration{
+		PackageName: "api",
+		Generate: GenerateOptions{
+			Models: true,
+		},
+		OutputOptions: OutputOptions{
+			UserTemplates:    map[string]string{"typedef.tmpl": "//blah"},
+			UserTemplatesDir: "test_templates",
+		},
+	}
+
+	// Run our code generation:
+	code, err = Generate(swagger, 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"}`)
 
 	cannedResponse := &http.Response{
@@ -129,7 +189,6 @@ func TestExamplePetStoreParseFunction(t *testing.T) {
 }
 
 func TestExampleOpenAPICodeGeneration(t *testing.T) {
-
 	// Input vars for code generation:
 	packageName := "testswagger"
 	opts := Configuration{
@@ -231,7 +290,6 @@ func TestXGoTypeImport(t *testing.T) {
 
 	// Make sure the generated code is valid:
 	checkLint(t, "test.gen.go", []byte(code))
-
 }
 
 //go:embed test_spec.yaml

diff --git a/pkg/codegen/configuration.go b/pkg/codegen/configuration.go
@@ -68,14 +68,14 @@ type CompatibilityOptions struct {
 
 // OutputOptions are used to modify the output code in some way.
 type OutputOptions struct {
-	SkipFmt       bool              `yaml:"skip-fmt,omitempty"`       // Whether to skip go imports on the generated code
-	SkipPrune     bool              `yaml:"skip-prune,omitempty"`     // Whether to skip pruning unused components on the generated code
-	IncludeTags   []string          `yaml:"include-tags,omitempty"`   // Only include operations that have one of these tags. Ignored when empty.
-	ExcludeTags   []string          `yaml:"exclude-tags,omitempty"`   // Exclude operations that have one of these tags. Ignored when empty.
-	UserTemplates map[string]string `yaml:"user-templates,omitempty"` // Override built-in templates from user-provided files
-
-	ExcludeSchemas     []string `yaml:"exclude-schemas,omitempty"`      // Exclude from generation schemas with given names. Ignored when empty.
-	ResponseTypeSuffix string   `yaml:"response-type-suffix,omitempty"` // The suffix used for responses types
+	SkipFmt            bool              `yaml:"skip-fmt,omitempty"`             // Whether to skip go imports on the generated code
+	SkipPrune          bool              `yaml:"skip-prune,omitempty"`           // Whether to skip pruning unused components on the generated code
+	IncludeTags        []string          `yaml:"include-tags,omitempty"`         // Only include operations that have one of these tags. Ignored when empty.
+	ExcludeTags        []string          `yaml:"exclude-tags,omitempty"`         // Exclude operations that have one of these tags. Ignored when empty.
+	UserTemplates      map[string]string `yaml:"user-templates,omitempty"`       // Override built-in templates from user-provided files
+	UserTemplatesDir   string            `yaml:"user-templates-dir,omitempty"`   // Override built-in templates from user-provided files by specifying a directory path
+	ExcludeSchemas     []string          `yaml:"exclude-schemas,omitempty"`      // Exclude from generation schemas with given names. Ignored when empty.
+	ResponseTypeSuffix string            `yaml:"response-type-suffix,omitempty"` // The suffix used for responses types
 }
 
 // UpdateDefaults sets reasonable default values for unset fields in Configuration

diff --git a/pkg/codegen/test_templates/chi/chi-interface.tmpl b/pkg/codegen/test_templates/chi/chi-interface.tmpl
@@ -0,0 +1 @@
+//__placeholder__