fixed custom template loading issue.

Author: HugoMario

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

@@ -11,20 +11,41 @@ public class CodegenTemplateLoader extends FileTemplateLoader {
     private String templateDir;
     private String customTemplateDir;
 
-    public CodegenTemplateLoader(String templateDir, String customTemplateDir, String suffix) {
-        super(customTemplateDir != null ? customTemplateDir : templateDir, suffix);
+    public CodegenTemplateLoader(String templateDir, String suffix) {
+        super(templateDir, suffix);
         this.templateDir = templateDir;
-        this.customTemplateDir = customTemplateDir;
     }
 
     public URL getResource(String location) throws IOException {
-        File file = new File(location);
+        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);
         if (file.exists()) {
             return file.toURI().toURL();
         }
-        location = location.replace(this.customTemplateDir, StringUtils.EMPTY);
-        location = this.templateDir + location;
-
         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;
+    }
+
+    public void setCustomTemplateDir(String customTemplateDir) {
+        this.customTemplateDir = customTemplateDir;
+    }
+
+    public CodegenTemplateLoader customTemplateDir(String customTemplateDir) {
+        this.customTemplateDir = customTemplateDir;
+        return this;
+    }
 }

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

@@ -31,19 +31,13 @@ public String getName() {
     }
 
     private com.github.jknack.handlebars.Template getHandlebars(String templateFile) throws IOException {
-        final boolean needFileTemplateLoader = StringUtils.isNotBlank(config.customTemplateDir());
-        final boolean fileExist = new File(templateFile).exists();
         templateFile = templateFile.replace(".mustache", StringUtils.EMPTY).replace("\\", "/");
         final String templateDir = config.templateDir().replace("\\", "/");
         final TemplateLoader templateLoader;
-        if (needFileTemplateLoader && fileExist) {
-            String customTemplateDir = config.customTemplateDir().replace("\\", "/");
-            templateFile = resolveTemplateFile(customTemplateDir, templateFile);
-            templateLoader = new CodegenTemplateLoader("/" + templateDir, customTemplateDir, ".mustache");
-        } else {
-            templateFile = resolveTemplateFile(templateDir, templateFile);
-            templateLoader = new ClassPathTemplateLoader("/" + templateDir, ".mustache");
-        }
+        String customTemplateDir = config.customTemplateDir() != null ? config.customTemplateDir().replace("\\", "/") : null;
+        templateFile = resolveTemplateFile(templateDir, templateFile);
+        templateLoader = new CodegenTemplateLoader("/" + templateDir, ".mustache")
+                .customTemplateDir(customTemplateDir);
         final Handlebars handlebars = new Handlebars(templateLoader);
         handlebars.prettyPrint(true);
         config.addHandlebarHelpers(handlebars);