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

#6685
baomidou · Feb 11, 2025 · d4ec84c · d4ec84c
1 parent 62dd02e
commit d4ec84c
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 6 deletions.
diff --git a/...lus-generator/src/main/java/com/baomidou/mybatisplus/generator/config/builder/Entity.java b/...lus-generator/src/main/java/com/baomidou/mybatisplus/generator/config/builder/Entity.java
@@ -53,6 +53,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /**
@@ -283,6 +284,21 @@ private Entity() {
      */
     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>
      * 父类 Class 反射属性转换为公共字段
@@ -418,15 +434,16 @@ public Map<String, Object> renderData(@NotNull TableInfo tableInfo) {
             tableInfo.getFields().forEach(tableField -> {
                 List<AnnotationAttributes> annotationAttributes = tableFieldAnnotationHandler.handle(tableInfo, tableField);
                 if (annotationAttributes != null && !annotationAttributes.isEmpty()) {
-                    tableField.addAnnotationAttributesList(annotationAttributes);
+                    tableField.addAnnotationAttributesList(annotationAttributes, annotationAttributesFunction);
                     annotationAttributes.forEach(attributes -> importPackages.addAll(attributes.getImportPackages()));
                 }
             });
         }
         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);
         return data;
     }
@@ -841,6 +858,29 @@ public Builder fieldUseJavaDoc(boolean fieldUseJavaDoc) {
             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() {
             String superClass = this.entity.superClass;

diff --git a/...plus-generator/src/main/java/com/baomidou/mybatisplus/generator/config/po/TableField.java b/...plus-generator/src/main/java/com/baomidou/mybatisplus/generator/config/po/TableField.java
@@ -35,6 +35,7 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /**
@@ -139,6 +140,13 @@ public class TableField {
      */
     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 Entity getEntity() {
      * @since 3.5.10
      */
     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.annotationAttributesFunction = annotationAttributesFunction;
     }
 
     /**
@@ -394,13 +411,13 @@ public void addAnnotationAttributesList(@NotNull AnnotationAttributes annotation
     }
 
     /**
-     * 获取字段注解属性(按{@link AnnotationAttributes#getDisplayName()}长度进行升序)
+     * 获取字段注解属性(默认按{@link AnnotationAttributes#getDisplayName()}长度进行升序)
      *
      * @return 字段注解属性
      * @since 3.5.10
      */
     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()))
             .collect(Collectors.toList());
     }