优化代码生成器

mybatis-plus/src/main/java/com/baomidou/mybatisplus/generator/AutoGenerator.java

@@ -35,7 +35,6 @@ import org.apache.velocity.app.Velocity;
 import org.apache.velocity.app.VelocityEngine;
 
 import com.baomidou.mybatisplus.generator.config.ConstVal;
-import com.baomidou.mybatisplus.generator.config.StrategyConfig;
 import com.baomidou.mybatisplus.generator.config.TemplateConfig;
 import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
 import com.baomidou.mybatisplus.generator.config.po.TableInfo;
@@ -118,20 +117,27 @@ public class AutoGenerator extends AbstractGenerator {
 				cfg.initMap();
 				ctx.put("cfg", cfg.getMap());
 			}
+			/* ---------- 添加导入包 ---------- */
+			if (!tableInfo.getEntityName().toLowerCase().equals(tableInfo.getName().toLowerCase())) {
+				// 表注解
+				tableInfo.setImportPackages("com.baomidou.mybatisplus.annotations.TableName");
+				ctx.put("tabeAnnotation", true);
+			}
+			if (StringUtils.isNotEmpty(config.getSuperEntityClass())) {
+				// 父实体
+				tableInfo.setImportPackages(config.getSuperEntityClass());
+			}
 			ctx.put("package", packageInfo);
 			ctx.put("author", config.getGlobalConfig().getAuthor());
+			ctx.put("activeRecord", config.getGlobalConfig().isActiveRecord());
 			ctx.put("date", date);
 			ctx.put("table", tableInfo);
-			ctx.put("dbColumnUnderline", StrategyConfig.DB_COLUMN_UNDERLINE);
-			ctx.put("activeRecord", config.getGlobalConfig().isActiveRecord());
 			ctx.put("enableCache", config.getGlobalConfig().isEnableCache());
 			ctx.put("baseResultMap", config.getGlobalConfig().isBaseResultMap());
 			ctx.put("baseColumnList", config.getGlobalConfig().isBaseColumnList());
 			ctx.put("entity", tableInfo.getEntityName());
 			ctx.put("entityColumnConstant", config.getStrategyConfig().isEntityColumnConstant());
 			ctx.put("entityBuliderModel", config.getStrategyConfig().isEntityBuliderModel());
-			ctx.put("tabeAnnotation", !tableInfo.getEntityName().toLowerCase().equals(tableInfo.getName().toLowerCase()));
-			ctx.put("superEntityClassPackage", config.getSuperEntityClass());
 			ctx.put("superEntityClass", superEntityClass);
 			ctx.put("superMapperClassPackage", config.getSuperMapperClass());
 			ctx.put("superMapperClass", superMapperClass);

mybatis-plus/src/main/java/com/baomidou/mybatisplus/generator/config/builder/ConfigBuilder.java

@@ -35,6 +35,7 @@ import com.baomidou.mybatisplus.generator.config.StrategyConfig;
 import com.baomidou.mybatisplus.generator.config.TemplateConfig;
 import com.baomidou.mybatisplus.generator.config.po.TableField;
 import com.baomidou.mybatisplus.generator.config.po.TableInfo;
+import com.baomidou.mybatisplus.generator.config.rules.DbColumnType;
 import com.baomidou.mybatisplus.generator.config.rules.DbType;
 import com.baomidou.mybatisplus.generator.config.rules.NamingStrategy;
 import com.baomidou.mybatisplus.generator.config.rules.QuerySQL;
@@ -433,7 +434,7 @@ public class ConfigBuilder {
 			}
 			field.setType(results.getString(querySQL.getFieldType()));
 			field.setPropertyName(processName(field.getName(), strategy));
-			field.setPropertyType(processFiledType(field.getType()));
+			field.setColumnType(processFiledType(field.getType()));
 			field.setComment(results.getString(querySQL.getFieldComment()));
 			fieldList.add(field);
 		}
@@ -481,7 +482,7 @@ public class ConfigBuilder {
 	 *
 	 * @return 转换成JAVA包装类型
 	 */
-	private String processFiledType(String type) {
+	private DbColumnType processFiledType(String type) {
 		if (QuerySQL.MYSQL == querySQL) {
 			return processMySqlType(type);
 		} else if (QuerySQL.ORACLE == querySQL) {
@@ -491,7 +492,7 @@ public class ConfigBuilder {
 		} else if (QuerySQL.POSTGRE_SQL == querySQL) {
 			return processPostgreSQL(type);
 		}
-		return null;
+		return DbColumnType.STRING;
 	}
 
 	/**
@@ -534,32 +535,32 @@ public class ConfigBuilder {
 	 *            字段类型
 	 * @return JAVA类型
 	 */
-	private String processMySqlType(String type) {
+	private DbColumnType processMySqlType(String type) {
 		String t = type.toLowerCase();
 		if (t.contains("char") || t.contains("text")) {
-			return "String";
+			return DbColumnType.STRING;
 		} else if (t.contains("bigint")) {
-			return "Long";
+			return DbColumnType.LONG;
 		} else if (t.contains("int")) {
-			return "Integer";
+			return DbColumnType.INTEGER;
 		} else if (t.contains("date") || t.contains("time") || t.contains("year")) {
-			return "java.util.Date";
+			return DbColumnType.DATE;
 		} else if (t.contains("text")) {
-			return "String";
+			return DbColumnType.STRING;
 		} else if (t.contains("bit")) {
-			return "Boolean";
+			return DbColumnType.BOOLEAN;
 		} else if (t.contains("decimal")) {
-			return "java.math.BigDecimal";
+			return DbColumnType.BIG_DECIMAL;
 		} else if (t.contains("blob")) {
-			return "byte[]";
+			return DbColumnType.BYTE_ARRAY;
 		} else if (t.contains("float")) {
-			return "Float";
+			return DbColumnType.FLOAT;
 		} else if (t.contains("double")) {
-			return "Double";
+			return DbColumnType.DOUBLE;
 		} else if (t.contains("json") || t.contains("enum")) {
-			return "String";
+			return DbColumnType.STRING;
 		}
-		return "String";
+		return DbColumnType.STRING;
 	}
 
 	/**
@@ -571,27 +572,27 @@ public class ConfigBuilder {
 	 *            字段类型
 	 * @return JAVA类型
 	 */
-	private String processOracleType(String type) {
+	private DbColumnType processOracleType(String type) {
 		String t = type.toUpperCase();
 		if (t.contains("CHAR")) {
-			return "String";
+			return DbColumnType.STRING;
 		} else if (t.contains("DATE") || t.contains("TIMESTAMP")) {
-			return "java.util.Date";
+			return DbColumnType.DATE;
 		} else if (t.contains("NUMBER")) {
 			if (t.matches("NUMBER\\(+\\d{1}+\\)")) {
-				return "Integer";
+				return DbColumnType.INTEGER;
 			} else if (t.matches("NUMBER\\(+\\d{2}+\\)")) {
-				return "Long";
+				return DbColumnType.LONG;
 			}
-			return "Double";
+			return DbColumnType.DOUBLE;
 		} else if (t.contains("FLOAT")) {
-			return "Float";
+			return DbColumnType.FLOAT;
 		} else if (t.contains("BLOB")) {
-			return "Object";
+			return DbColumnType.OBJECT;
 		} else if (t.contains("RAW")) {
-			return "byte[]";
+			return DbColumnType.BYTE_ARRAY;
 		}
-		return "String";
+		return DbColumnType.STRING;
 	}
 
 	/**
@@ -602,30 +603,30 @@ public class ConfigBuilder {
 	 * @param type
 	 * @return
 	 */
-	private String processSQLServerType(String type) {
+	private DbColumnType processSQLServerType(String type) {
 		String t = type.toLowerCase();
 		if (t.contains("char") || t.contains("text") || t.contains("xml")) {
-			return "String";
+			return DbColumnType.STRING;
 		} else if (t.contains("bigint")) {
-			return "Long";
+			return DbColumnType.LONG;
 		} else if (t.contains("int")) {
-			return "Integer";
+			return DbColumnType.INTEGER;
 		} else if (t.contains("date") || t.contains("time")) {
-			return "java.util.Date";
+			return DbColumnType.DATE;
 		} else if (t.contains("text")) {
-			return "String";
+			return DbColumnType.STRING;
 		} else if (t.contains("bit")) {
-			return "Boolean";
+			return DbColumnType.BOOLEAN;
 		}else if (t.contains("decimal")||t.contains("numeric")){
-			return "Double";
+			return DbColumnType.DOUBLE;
 		}else if (t.contains("money")) {
-			return "java.math.BigDecimal";
+			return DbColumnType.BIG_DECIMAL;
 		} else if (t.contains("binary")||t.contains("image")) {
-			return "byte[]";
+			return DbColumnType.BYTE_ARRAY;
 		} else if (t.contains("float")||t.contains("real")) {
-			return "Float";
+			return DbColumnType.FLOAT;
 		}
-		return "String";
+		return DbColumnType.STRING;
 	}
 
 
@@ -638,32 +639,32 @@ public class ConfigBuilder {
 	 *            字段类型
 	 * @return JAVA类型
 	 */
-	private String processPostgreSQL(String type) {
+	private DbColumnType processPostgreSQL(String type) {
 		String t = type.toLowerCase();
 		if (t.contains("char") || t.contains("text")) {
-			return "String";
+			return DbColumnType.STRING;
 		} else if (t.contains("bigint")) {
-			return "Long";
+			return DbColumnType.LONG;
 		} else if (t.contains("int")) {
-			return "Integer";
+			return DbColumnType.INTEGER;
 		} else if (t.contains("date") || t.contains("time") || t.contains("year")) {
-			return "java.util.Date";
+			return DbColumnType.DATE;
 		} else if (t.contains("text")) {
-			return "String";
+			return DbColumnType.STRING;
 		} else if (t.contains("bit")) {
-			return "Boolean";
+			return DbColumnType.BOOLEAN;
 		} else if (t.contains("decimal")) {
-			return "java.math.BigDecimal";
+			return DbColumnType.BIG_DECIMAL;
 		} else if (t.contains("blob")) {
-			return "byte[]";
+			return DbColumnType.BYTE_ARRAY;
 		} else if (t.contains("float")) {
-			return "Float";
+			return DbColumnType.FLOAT;
 		} else if (t.contains("double")) {
-			return "Double";
+			return DbColumnType.DOUBLE;
 		} else if (t.contains("json") || t.contains("enum")) {
-			return "String";
+			return DbColumnType.STRING;
 		}
-		return "String";
+		return DbColumnType.STRING;
 	}
 
 	/**

mybatis-plus/src/main/java/com/baomidou/mybatisplus/generator/config/po/TableField.java

@@ -16,11 +16,13 @@
 package com.baomidou.mybatisplus.generator.config.po;
 
 import com.baomidou.mybatisplus.generator.config.StrategyConfig;
+import com.baomidou.mybatisplus.generator.config.rules.DbColumnType;
 
 /**
  * <p>
  * 表字段信息
  * </p>
+ * 
  * @author YangHu
  * @since 2016-12-03
  */
@@ -29,7 +31,7 @@ public class TableField {
 	private String name;
 	private String type;
 	private String propertyName;
-	private String propertyType;
+	private DbColumnType columnType;
 	private String comment;
 
 	public boolean isKeyFlag() {
@@ -64,12 +66,19 @@ public class TableField {
 		this.propertyName = propertyName;
 	}
 
-	public String getPropertyType() {
-		return propertyType;
+	public DbColumnType getColumnType() {
+		return columnType;
+	}
+
+	public void setColumnType(DbColumnType columnType) {
+		this.columnType = columnType;
 	}
 
-	public void setPropertyType(String propertyType) {
-		this.propertyType = propertyType;
+	public String getPropertyType() {
+		if (null != columnType) {
+			return columnType.getType();
+		}
+		return null;
 	}
 
 	public String getComment() {

mybatis-plus/src/main/java/com/baomidou/mybatisplus/generator/config/po/TableInfo.java

@@ -15,8 +15,13 @@
  */
 package com.baomidou.mybatisplus.generator.config.po;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
+import com.baomidou.mybatisplus.toolkit.CollectionUtils;
 import com.baomidou.mybatisplus.toolkit.StringUtils;
 
 /**
@@ -40,8 +45,8 @@ public class TableInfo {
 	private String controllerName;
 
 	private List<TableField> fields;
+	private List<String> importPackages = new ArrayList<String>();
 	private String fieldNames;
-	private boolean hasDate;
 
 	public String getName() {
 		return name;
@@ -119,7 +124,33 @@ public class TableInfo {
 	}
 
 	public void setFields(List<TableField> fields) {
-		this.fields = fields;
+		if (CollectionUtils.isNotEmpty(fields)) {
+			this.fields = fields;
+			// 收集导入包信息
+			Set<String> pkgSet = new HashSet<String>();
+			for (TableField field : fields) {
+				if (null != field.getColumnType() && null != field.getColumnType().getPkg()) {
+					pkgSet.add(field.getColumnType().getPkg());
+				}
+				if (field.isKeyFlag() && field.isConvert()) {
+					pkgSet.add("com.baomidou.mybatisplus.annotations.TableId");
+				}
+				if (field.isConvert()) {
+					pkgSet.add("com.baomidou.mybatisplus.annotations.TableField");
+				}
+			}
+			if (!pkgSet.isEmpty()) {
+				this.importPackages = Arrays.asList(pkgSet.toArray(new String[] {}));
+			}
+		}
+	}
+
+	public List<String> getImportPackages() {
+		return importPackages;
+	}
+
+	public void setImportPackages(String pkg) {
+		importPackages.add(pkg);
 	}
 
 	/**
@@ -143,21 +174,6 @@ public class TableInfo {
 		return fieldNames;
 	}
 
-	/**
-	 * 判断字段中是否包含日期类型
-	 *
-	 * @return 是否
-	 */
-	public boolean isHasDate() {
-		for (TableField fieldInfo : fields) {
-			if (fieldInfo.getPropertyType().equals("Date")) {
-				hasDate = true;
-				break;
-			}
-		}
-		return hasDate;
-	}
-
 	/**
 	 * mapper xml中的字字段添加as
 	 *

mybatis-plus/src/main/resources/template/controller.java.vm

@@ -8,7 +8,7 @@ import ${superControllerClassPackage};
 
 /**
  * <p>
- * ${table.comment}  前端控制器
+ * ${table.comment} 前端控制器
  * </p>
  *
  * @author ${author}

mybatis-plus/src/main/resources/template/entity.java.vm

@@ -1,19 +1,14 @@
 package ${package.Entity};
 
-#if(${superEntityClassPackage})
-import ${superEntityClassPackage};
-#end
 #if(${activeRecord})
 import com.baomidou.mybatisplus.activerecord.Model;
 #end
-#if(!${dbColumnUnderline})
-import com.baomidou.mybatisplus.annotations.TableId;
-import com.baomidou.mybatisplus.annotations.TableField;
-#end
-#if(${tabeAnnotation})
-import com.baomidou.mybatisplus.annotations.TableName;
+#foreach($pkg in ${table.importPackages})
+import ${pkg};
 #end
 import java.io.Serializable;
+
+
 /**
  * <p>
  * ${table.comment}
@@ -46,12 +41,10 @@ public class ${entity} implements Serializable {
 	 * ${field.comment}
 	 */
 #end
+#if(${field.convert})
 #if(${field.keyFlag})
-	@TableId#if(${field.convert})(value="${field.name}")
-#end
-
+	@TableId(value="${field.name}")
 #else
-#if(${field.convert})
 	@TableField(value="${field.name}")
 #end
 #end