001package cn.mybatis.mp.generator.template;
002
003import cn.mybatis.mp.generator.config.GeneratorConfig;
004import cn.mybatis.mp.generator.database.meta.EntityInfo;
005import cn.mybatis.mp.generator.util.GeneratorUtil;
006
007import java.time.LocalDate;
008import java.util.HashMap;
009import java.util.Map;
010
011public class EntityTemplateBuilder extends AbstractTemplateBuilder {
012
013    public EntityTemplateBuilder(GeneratorConfig generatorConfig, EntityInfo entityInfo) {
014        super(generatorConfig, entityInfo);
015    }
016
017    @Override
018    public String targetFilePath() {
019        return generatorConfig.getBaseFilePath() + "/" + (entityInfo.getEntityPackage() + "." + entityInfo.getName()).replaceAll("\\.", "/") + ".java";
020    }
021
022    @Override
023    public String templateFilePath() {
024        return generatorConfig.getTemplateRootPath() + "/entity";
025    }
026
027    @Override
028    public Map<String, Object> contextData() {
029        Map<String, Object> data = new HashMap<>();
030        data.put("imports", GeneratorUtil.buildEntityImports(generatorConfig, entityInfo));
031        if (generatorConfig.getEntityConfig().getSuperClass() != null) {
032            int dotIndex = generatorConfig.getEntityConfig().getSuperClass().lastIndexOf(".");
033            String superName;
034            if (dotIndex > 0) {
035                superName = generatorConfig.getEntityConfig().getSuperClass().substring(dotIndex + 1);
036            } else {
037                superName = generatorConfig.getEntityConfig().getSuperClass();
038            }
039            data.put("superExtend", "extends " + superName);
040        } else {
041            data.put("superExtend", "");
042        }
043        data.put("package", entityInfo.getEntityPackage());
044        data.put("date", LocalDate.now().toString());
045        data.put("author", generatorConfig.getAuthor());
046        data.put("entityInfo", entityInfo);
047        data.put("entityConfig", generatorConfig.getEntityConfig());
048        return data;
049    }
050}