瀏覽代碼

开放实体导包和字段注解处理(支持排序和动态添加).

https://github.com/baomidou/mybatis-plus/issues/6685
nieqiurong 4 月之前
父節點
當前提交
d4ec84c277

+ 44 - 4
mybatis-plus-generator/src/main/java/com/baomidou/mybatisplus/generator/config/builder/Entity.java

@@ -53,6 +53,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Optional;
 import java.util.Set;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Collectors;
 
 
 /**
 /**
@@ -283,6 +284,21 @@ public class Entity implements ITemplate {
      */
      */
     private ITableFieldAnnotationHandler tableFieldAnnotationHandler = new DefaultTableFieldAnnotationHandler();
     private ITableFieldAnnotationHandler tableFieldAnnotationHandler = new DefaultTableFieldAnnotationHandler();
 
 
+    /**
+     * 导包处理方法
+     *
+     * @since 3.5.10.2
+     */
+    private Function<Set<String>, List<String>> importPackageFunction;
+
+    /**
+     * 处理类注解方法 (含类与字段)
+     *
+     * @since 3.5.10.2
+     */
+    private Function<List<? extends AnnotationAttributes>, List<AnnotationAttributes>> annotationAttributesFunction;
+
+
     /**
     /**
      * <p>
      * <p>
      * 父类 Class 反射属性转换为公共字段
      * 父类 Class 反射属性转换为公共字段
@@ -418,15 +434,16 @@ public class Entity implements ITemplate {
             tableInfo.getFields().forEach(tableField -> {
             tableInfo.getFields().forEach(tableField -> {
                 List<AnnotationAttributes> annotationAttributes = tableFieldAnnotationHandler.handle(tableInfo, tableField);
                 List<AnnotationAttributes> annotationAttributes = tableFieldAnnotationHandler.handle(tableInfo, tableField);
                 if (annotationAttributes != null && !annotationAttributes.isEmpty()) {
                 if (annotationAttributes != null && !annotationAttributes.isEmpty()) {
-                    tableField.addAnnotationAttributesList(annotationAttributes);
+                    tableField.addAnnotationAttributesList(annotationAttributes, annotationAttributesFunction);
                     annotationAttributes.forEach(attributes -> importPackages.addAll(attributes.getImportPackages()));
                     annotationAttributes.forEach(attributes -> importPackages.addAll(attributes.getImportPackages()));
                 }
                 }
             });
             });
         }
         }
         data.put("entityFieldUseJavaDoc", fieldUseJavaDoc);
         data.put("entityFieldUseJavaDoc", fieldUseJavaDoc);
-        data.put("entityClassAnnotations", classAnnotationAttributes.stream()
-            .sorted(Comparator.comparingInt(s -> s.getDisplayName().length())).collect(Collectors.toList()));
-        data.put("importEntityPackages", importPackages.stream().sorted().collect(Collectors.toList()));
+        data.put("entityClassAnnotations", annotationAttributesFunction != null ? annotationAttributesFunction.apply(classAnnotationAttributes) :
+            classAnnotationAttributes.stream().sorted(Comparator.comparingInt(s -> s.getDisplayName().length())).collect(Collectors.toList()));
+        data.put("importEntityPackages", importPackageFunction != null ? importPackageFunction.apply(importPackages) :
+            importPackages.stream().sorted().collect(Collectors.toList()));
         data.put("entityToString", this.toString);
         data.put("entityToString", this.toString);
         return data;
         return data;
     }
     }
@@ -841,6 +858,29 @@ public class Entity implements ITemplate {
             return this;
             return this;
         }
         }
 
 
+        /**
+         * 导包处理方法
+         *
+         * @param importPackageFunction 导包处理
+         * @return this
+         * @since 3.5.10.2
+         */
+        public Builder importPackageFunction(Function<Set<String>, List<String>> importPackageFunction) {
+            this.entity.importPackageFunction = importPackageFunction;
+            return this;
+        }
+
+        /**
+         * 注解处理方法 (含类与字段)
+         *
+         * @param annotationAttributesFunction 注解处理
+         * @return this
+         * @since 3.5.10.2
+         */
+        public Builder annotationAttributesFunction(Function<List<? extends AnnotationAttributes>, List<AnnotationAttributes>> annotationAttributesFunction) {
+            this.entity.annotationAttributesFunction = annotationAttributesFunction;
+            return this;
+        }
 
 
         public Entity get() {
         public Entity get() {
             String superClass = this.entity.superClass;
             String superClass = this.entity.superClass;

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

@@ -35,6 +35,7 @@ import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.Comparator;
 import java.util.List;
 import java.util.List;
 import java.util.Map;
 import java.util.Map;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Collectors;
 
 
 /**
 /**
@@ -139,6 +140,13 @@ public class TableField {
      */
      */
     private final List<AnnotationAttributes> annotationAttributesList = new ArrayList<>();
     private final List<AnnotationAttributes> annotationAttributesList = new ArrayList<>();
 
 
+    /**
+     * 字段注解处理
+     *
+     * @since 3.5.10.2
+     */
+    private Function<List<? extends AnnotationAttributes>, List<AnnotationAttributes>> annotationAttributesFunction;
+
     /**
     /**
      * 构造方法
      * 构造方法
      *
      *
@@ -380,7 +388,16 @@ public class TableField {
      * @since 3.5.10
      * @since 3.5.10
      */
      */
     public void addAnnotationAttributesList(@NotNull List<AnnotationAttributes> annotationAttributesList) {
     public void addAnnotationAttributesList(@NotNull List<AnnotationAttributes> annotationAttributesList) {
+        this.addAnnotationAttributesList(annotationAttributesList, null);
+    }
+
+    /**
+     * @param annotationAttributesList 注解属性集合
+     * @since 3.5.10.2
+     */
+    public void addAnnotationAttributesList(@NotNull List<AnnotationAttributes> annotationAttributesList, Function<List<? extends AnnotationAttributes>, List<AnnotationAttributes>> annotationAttributesFunction) {
         this.annotationAttributesList.addAll(annotationAttributesList);
         this.annotationAttributesList.addAll(annotationAttributesList);
+        this.annotationAttributesFunction = annotationAttributesFunction;
     }
     }
 
 
     /**
     /**
@@ -394,13 +411,13 @@ public class TableField {
     }
     }
 
 
     /**
     /**
-     * 获取字段注解属性(按{@link AnnotationAttributes#getDisplayName()}长度进行升序)
+     * 获取字段注解属性(默认按{@link AnnotationAttributes#getDisplayName()}长度进行升序)
      *
      *
      * @return 字段注解属性
      * @return 字段注解属性
      * @since 3.5.10
      * @since 3.5.10
      */
      */
     public List<AnnotationAttributes> getAnnotationAttributesList() {
     public List<AnnotationAttributes> getAnnotationAttributesList() {
-        return this.annotationAttributesList.stream()
+        return annotationAttributesFunction != null ? annotationAttributesFunction.apply(this.annotationAttributesList) : this.annotationAttributesList.stream()
             .sorted(Comparator.comparingInt(s -> s.getDisplayName().length()))
             .sorted(Comparator.comparingInt(s -> s.getDisplayName().length()))
             .collect(Collectors.toList());
             .collect(Collectors.toList());
     }
     }