CustomTemplateEngine.java 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  1. package jnpf.base.util.custom;
  2. import com.baomidou.mybatisplus.core.toolkit.Constants;
  3. import com.baomidou.mybatisplus.generator.config.ConstVal;
  4. import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
  5. import com.baomidou.mybatisplus.generator.engine.AbstractTemplateEngine;
  6. import org.apache.velocity.Template;
  7. import org.apache.velocity.VelocityContext;
  8. import org.apache.velocity.app.VelocityEngine;
  9. import org.jetbrains.annotations.NotNull;
  10. import java.io.*;
  11. import java.util.Map;
  12. import java.util.Properties;
  13. public class CustomTemplateEngine extends AbstractTemplateEngine {
  14. private static final String DOT_VM = ".vm";
  15. private VelocityEngine velocityEngine;
  16. private Map<String, Object> customParams;
  17. private String path;
  18. public CustomTemplateEngine(String path) {
  19. this.path = path;
  20. }
  21. public CustomTemplateEngine(Map<String, Object> customParams, String path) {
  22. this.customParams = customParams;
  23. this.path = path;
  24. }
  25. @Override
  26. public CustomTemplateEngine init(ConfigBuilder configBuilder) {
  27. super.setConfigBuilder(configBuilder);
  28. if (null == this.velocityEngine) {
  29. Properties p = new Properties();
  30. p.setProperty(VelocityEngine.FILE_RESOURCE_LOADER_PATH, path);
  31. p.setProperty("ISO-8859-1", Constants.UTF_8);
  32. p.setProperty("output.encoding", Constants.UTF_8);
  33. this.velocityEngine = new VelocityEngine(p);
  34. }
  35. return this;
  36. }
  37. @Override
  38. public String writer(@NotNull Map<String, Object> objectMap, @NotNull String templateName, @NotNull String templateString) throws Exception {
  39. StringWriter writer = new StringWriter();
  40. this.velocityEngine.evaluate(new VelocityContext(objectMap), writer, templateName, templateString);
  41. return writer.toString();
  42. }
  43. @Override
  44. public void writer(Map<String, Object> objectMap, String templatePath, File outputFile) throws Exception {
  45. if (templatePath != null && !"".equals(templatePath.trim())) {
  46. Template template = this.velocityEngine.getTemplate(templatePath, ConstVal.UTF8);
  47. FileOutputStream fos = new FileOutputStream(outputFile);
  48. BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(fos, ConstVal.UTF8));
  49. if (customParams != null) {
  50. objectMap.putAll(customParams);
  51. }
  52. template.merge(new VelocityContext(objectMap), writer);
  53. writer.close();
  54. }
  55. }
  56. @Override
  57. public String templateFilePath(String filePath) {
  58. if (null != filePath && !filePath.contains(".vm")) {
  59. StringBuilder fp = new StringBuilder();
  60. fp.append(filePath).append(".vm");
  61. return fp.toString();
  62. } else {
  63. return filePath;
  64. }
  65. }
  66. }