001package cn.mybatis.mp.generator.template.engine;
002
003import cn.mybatis.mp.generator.template.ITemplateBuilder;
004import freemarker.core.PlainTextOutputFormat;
005import freemarker.template.Configuration;
006import freemarker.template.Template;
007import freemarker.template.TemplateModelException;
008
009import java.io.File;
010import java.io.FileOutputStream;
011import java.io.IOException;
012import java.io.OutputStreamWriter;
013import java.nio.charset.StandardCharsets;
014
015public class FreemarkerTemplateEngine implements TemplateEngine {
016
017    private final Configuration configuration;
018
019    public FreemarkerTemplateEngine() {
020        configuration = new Configuration(Configuration.DEFAULT_INCOMPATIBLE_IMPROVEMENTS);
021        configuration.setDefaultEncoding("UTF-8");
022
023        configuration.setOutputFormat(PlainTextOutputFormat.INSTANCE);
024        configuration.setClassForTemplateLoading(FreemarkerTemplateEngine.class, "/");
025        try {
026            configuration.setSharedVariable("util", new EngineUtil());
027        } catch (TemplateModelException e) {
028            throw new RuntimeException(e);
029        }
030    }
031
032    @Override
033    public void render(ITemplateBuilder templateBuilder) {
034        try {
035            Template template = configuration.getTemplate(templateBuilder.templateFilePath() + ".ftl");
036
037            File outputFile = new File(templateBuilder.targetFilePath());
038            if (!outputFile.exists()) {
039                outputFile.getParentFile().mkdirs();
040                outputFile.createNewFile();
041            }
042
043            try (FileOutputStream fileOutputStream = new FileOutputStream(outputFile)) {
044                template.process(templateBuilder.contextData(), new OutputStreamWriter(fileOutputStream, StandardCharsets.UTF_8));
045            } catch (Exception e) {
046                outputFile.delete();
047                throw new RuntimeException(e);
048            }
049        } catch (IOException e) {
050            throw new RuntimeException(e);
051        }
052    }
053}