Merge pull request #11689 from swagger-api/issue-11621

Fixed custom templates issue

Author: HugoMario

modules/swagger-codegen/src/main/java/io/swagger/codegen/v3/templates/CodegenTemplateLoader.java

@@ -1,41 +1,56 @@
 package io.swagger.codegen.v3.templates;
 
-import com.github.jknack.handlebars.io.FileTemplateLoader;
-import org.apache.commons.lang3.StringUtils;
+import com.github.jknack.handlebars.io.URLTemplateLoader;
 
 import java.io.File;
 import java.io.IOException;
 import java.net.URL;
 
-public class CodegenTemplateLoader extends FileTemplateLoader {
+public class CodegenTemplateLoader extends URLTemplateLoader {
     private String templateDir;
     private String customTemplateDir;
 
-    public CodegenTemplateLoader(String templateDir, String suffix) {
-        super(templateDir, suffix);
-        this.templateDir = templateDir;
+    public CodegenTemplateLoader() {
+        setSuffix(".mustache");
     }
 
+    @Override
+    public String resolve(String uri) {
+        if (!uri.endsWith(this.getSuffix())) {
+            uri = uri + this.getSuffix();
+        }
+        File templateFile = new File(uri);
+        if (templateFile.exists()) {
+            return templateFile.toString();
+        }
+        templateFile = new File(this.getPrefix() + this.normalize(uri));
+        if (templateFile.exists()) {
+            return templateFile.toString();
+        }
+        if (this.customTemplateDir != null) {
+            templateFile = new File(this.customTemplateDir, this.normalize(uri));
+            if (templateFile.exists()) {
+                return templateFile.toString();
+            }
+        }
+        if (getClass().getResource(this.getPrefix() + this.normalize(uri)) != null) {
+            return this.getPrefix() + this.normalize(uri);
+        }
+        return this.templateDir + this.normalize(uri);
+    }
+
+    @Override
     public URL getResource(String location) throws IOException {
         if (this.customTemplateDir == null) {
             return this.getClass().getResource(location);
         }
-        String templateFile = resolveTemplateFile(this.templateDir, location);
-        templateFile = resolveTemplateFile(this.customTemplateDir, templateFile);
-        final File file = new File(this.customTemplateDir, templateFile);
+        final File file = new File(location);
         if (file.exists()) {
             return file.toURI().toURL();
         }
         return this.getClass().getResource(location);
     }
 
-    private String resolveTemplateFile(String templateDir, String templateFile) {
-        if (templateFile.startsWith(templateDir)) {
-            templateFile = StringUtils.replaceOnce(templateFile, templateDir, StringUtils.EMPTY);
-        }
-        return templateFile;
-    }
-
     public String getCustomTemplateDir() {
         return customTemplateDir;
     }
@@ -48,4 +63,20 @@ public CodegenTemplateLoader customTemplateDir(String customTemplateDir) {
         this.customTemplateDir = customTemplateDir;
         return this;
     }
+
+    public String getTemplateDir() {
+        return templateDir;
+    }
+
+    public void setTemplateDir(String templateDir) {
+        this.templateDir = this.getPrefix() + templateDir;
+        if (!this.templateDir.endsWith("/")) {
+            this.templateDir = this.templateDir + "/";
+        }
+    }
+
+    public CodegenTemplateLoader templateDir(String templateDir) {
+        setTemplateDir(templateDir);
+        return this;
+    }
 }

modules/swagger-codegen/src/main/java/io/swagger/codegen/v3/templates/HandlebarTemplateEngine.java

@@ -1,13 +1,10 @@
 package io.swagger.codegen.v3.templates;
 
 import com.github.jknack.handlebars.Handlebars;
-import com.github.jknack.handlebars.io.ClassPathTemplateLoader;
 import com.github.jknack.handlebars.io.TemplateLoader;
 import io.swagger.codegen.v3.CodegenConfig;
 import io.swagger.codegen.v3.CodegenConstants;
-import org.apache.commons.lang3.StringUtils;
 
-import java.io.File;
 import java.io.IOException;
 import java.util.Map;
 
@@ -31,23 +28,16 @@ public String getName() {
     }
 
     private com.github.jknack.handlebars.Template getHandlebars(String templateFile) throws IOException {
-        templateFile = templateFile.replace(".mustache", StringUtils.EMPTY).replace("\\", "/");
+        templateFile = templateFile.replace("\\", "/");
         final String templateDir = config.templateDir().replace("\\", "/");
         final TemplateLoader templateLoader;
         String customTemplateDir = config.customTemplateDir() != null ? config.customTemplateDir().replace("\\", "/") : null;
-        templateFile = resolveTemplateFile(templateDir, templateFile);
-        templateLoader = new CodegenTemplateLoader("/" + templateDir, ".mustache")
+        templateLoader = new CodegenTemplateLoader()
+                .templateDir(templateDir)
                 .customTemplateDir(customTemplateDir);
         final Handlebars handlebars = new Handlebars(templateLoader);
         handlebars.prettyPrint(true);
         config.addHandlebarHelpers(handlebars);
         return handlebars.compile(templateFile);
     }
-
-    private String resolveTemplateFile(String templateDir, String templateFile) {
-        if (templateFile.startsWith(templateDir)) {
-            templateFile = StringUtils.replaceOnce(templateFile, templateDir, StringUtils.EMPTY);
-        }
-        return templateFile;
-    }
 }