feat: 增加从实体类生成代码功能.

src/main/java/com/sjhy/plugin/actions/EasyCodeEntityAction.java

@@ -0,0 +1,53 @@
+package com.sjhy.plugin.actions;
+
+import com.google.common.collect.Lists;
+import com.intellij.openapi.actionSystem.AnAction;
+import com.intellij.openapi.actionSystem.AnActionEvent;
+import com.intellij.openapi.actionSystem.CommonDataKeys;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.*;
+import com.intellij.psi.javadoc.PsiDocComment;
+import com.intellij.psi.javadoc.PsiDocToken;
+import com.sjhy.plugin.entity.ColumnInfo;
+import com.sjhy.plugin.entity.TableInfo;
+import com.sjhy.plugin.tool.CacheDataUtils;
+import com.sjhy.plugin.ui.SelectSavePath;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+
+/**
+ * 从Java类生成代码菜单
+ */
+public class EasyCodeEntityAction extends AnAction {
+
+    private CacheDataUtils cacheDataUtils = CacheDataUtils.getInstance();
+
+    @Override
+    public void actionPerformed(@NotNull AnActionEvent event) {
+        Project project = event.getProject();
+        if (project == null) {
+            return;
+        }
+        PsiFile psiFile = event.getData(CommonDataKeys.PSI_FILE);
+        if (!(psiFile instanceof PsiJavaFile)) {
+            return;
+        }
+        PsiJavaFile psiJavaFile = (PsiJavaFile) psiFile;
+        PsiClass psiClass = Arrays.stream(psiJavaFile.getClasses())
+                .filter(o -> o.getModifierList() != null && o.getModifierList().hasModifierProperty("public"))
+                .findFirst()
+                .orElse(null);
+        if (psiClass == null) {
+            return;
+        }
+        cacheDataUtils.setSelectPsiClass(psiClass);
+        cacheDataUtils.setPsiClassList(Lists.newArrayList(psiClass));
+        new SelectSavePath(project, true).open();
+    }
+
+
+
+}

src/main/java/com/sjhy/plugin/entity/TableInfo.java

@@ -2,6 +2,7 @@ package com.sjhy.plugin.entity;
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.intellij.database.psi.DbTable;
+import com.intellij.psi.PsiClass;
 import lombok.Data;
 
 import java.util.List;
@@ -20,6 +21,13 @@ public class TableInfo {
      */
     @JsonIgnore
     private DbTable obj;
+
+    /**
+     * 原始对象（从实体生成）
+     */
+    @JsonIgnore
+    private PsiClass psiClassObj;
+
     /**
      * 表名（首字母大写）
      */

src/main/java/com/sjhy/plugin/service/CodeGenerateService.java

@@ -28,12 +28,13 @@ public interface CodeGenerateService {
 
     /**
      * 生成代码，并自动保存到对应位置，使用统一配置
-     *
+     *  @param templates     模板
-     * @param templates     模板
      * @param unifiedConfig 是否使用统一配置
      * @param title         是否显示提示
+     * @param entityMode
      */
-    void generateByUnifiedConfig(Collection<Template> templates, boolean unifiedConfig, boolean title);
+    void generateByUnifiedConfig(Collection<Template> templates, boolean unifiedConfig, boolean title,
+            boolean entityMode);
 
     /**
      * 生成代码

src/main/java/com/sjhy/plugin/service/TableInfoService.java

@@ -3,6 +3,7 @@ package com.sjhy.plugin.service;
 import com.intellij.database.psi.DbTable;
 import com.intellij.openapi.components.ServiceManager;
 import com.intellij.openapi.project.Project;
+import com.intellij.psi.PsiClass;
 import com.sjhy.plugin.entity.TableInfo;
 import com.sjhy.plugin.tool.CollectionUtil;
 import org.jetbrains.annotations.NotNull;
@@ -92,4 +93,19 @@ public interface TableInfoService {
      * @param tableInfo 表信息对象
      */
     void save(TableInfo tableInfo);
+
+    /**
+     * 获取表信息
+     */
+    TableInfo getTableInfoByPsiClass(PsiClass psiClass);
+
+    /**
+     * 从实体获取表信息包含配置
+     */
+    TableInfo getTableInfoAndConfigByPsiClass(PsiClass selectPsiClass);
+
+    /**
+     * 从实体获取表信息包含配置
+     */
+    List<TableInfo> getTableInfoAndConfigByPsiClass(List<PsiClass> selectPsiClass);
 }

src/main/java/com/sjhy/plugin/service/impl/CodeGenerateServiceImpl.java

@@ -16,9 +16,24 @@ import com.sjhy.plugin.entity.TableInfo;
 import com.sjhy.plugin.entity.Template;
 import com.sjhy.plugin.service.CodeGenerateService;
 import com.sjhy.plugin.service.TableInfoService;
-import com.sjhy.plugin.tool.*;
+import com.sjhy.plugin.tool.CacheDataUtils;
-
+import com.sjhy.plugin.tool.CloneUtils;
-import java.util.*;
+import com.sjhy.plugin.tool.CollectionUtil;
+import com.sjhy.plugin.tool.ExtraCodeGenerateUtils;
+import com.sjhy.plugin.tool.GlobalTool;
+import com.sjhy.plugin.tool.ModuleUtils;
+import com.sjhy.plugin.tool.NameUtils;
+import com.sjhy.plugin.tool.StringUtils;
+import com.sjhy.plugin.tool.TemplateUtils;
+import com.sjhy.plugin.tool.TimeUtils;
+import com.sjhy.plugin.tool.VelocityUtils;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeSet;
 
 /**
  * @author makejava
@@ -56,37 +71,45 @@ public class CodeGenerateServiceImpl implements CodeGenerateService {
 
     /**
      * 生成代码，并自动保存到对应位置，使用统一配置
-     *
+     *  @param templates     模板
-     * @param templates     模板
      * @param unifiedConfig 是否使用统一配置
      * @param title         是否显示提示
+     * @param entityMode
      */
     @Override
-    public void generateByUnifiedConfig(Collection<Template> templates, boolean unifiedConfig, boolean title) {
+    public void generateByUnifiedConfig(Collection<Template> templates, boolean unifiedConfig, boolean title,
+            boolean entityMode) {
         // 获取选中表信息
-        TableInfo selectedTableInfo = tableInfoService.getTableInfoAndConfig(cacheDataUtils.getSelectDbTable());
+        TableInfo selectedTableInfo;
-        // 获取所有选中的表信息
+        List<TableInfo> tableInfoList;
-        List<TableInfo> tableInfoList = tableInfoService.getTableInfoAndConfig(cacheDataUtils.getDbTableList());
+        if(!entityMode) {
+            selectedTableInfo = tableInfoService.getTableInfoAndConfig(cacheDataUtils.getSelectDbTable());
+            tableInfoList = tableInfoService.getTableInfoAndConfig(cacheDataUtils.getDbTableList());
+        } else {
+            selectedTableInfo = tableInfoService.getTableInfoAndConfigByPsiClass(cacheDataUtils.getSelectPsiClass());
+            tableInfoList = tableInfoService.getTableInfoAndConfigByPsiClass(cacheDataUtils.getPsiClassList());
+        }
         // 校验选中表的保存路径是否正确
         if (StringUtils.isEmpty(selectedTableInfo.getSavePath())) {
             Messages.showInfoMessage(selectedTableInfo.getObj().getName() + "表配置信息不正确，请尝试重新配置", MsgValue.TITLE_INFO);
             return;
         }
         // 将未配置的表进行配置覆盖
+        TableInfo finalSelectedTableInfo = selectedTableInfo;
         tableInfoList.forEach(tableInfo -> {
             if (StringUtils.isEmpty(tableInfo.getSavePath())) {
-                tableInfo.setSaveModelName(selectedTableInfo.getSaveModelName());
+                tableInfo.setSaveModelName(finalSelectedTableInfo.getSaveModelName());
-                tableInfo.setSavePackageName(selectedTableInfo.getSavePackageName());
+                tableInfo.setSavePackageName(finalSelectedTableInfo.getSavePackageName());
-                tableInfo.setSavePath(selectedTableInfo.getSavePath());
+                tableInfo.setSavePath(finalSelectedTableInfo.getSavePath());
                 tableInfoService.save(tableInfo);
             }
         });
         // 如果使用统一配置，直接全部覆盖
         if (unifiedConfig) {
             tableInfoList.forEach(tableInfo -> {
-                tableInfo.setSaveModelName(selectedTableInfo.getSaveModelName());
+                tableInfo.setSaveModelName(finalSelectedTableInfo.getSaveModelName());
-                tableInfo.setSavePackageName(selectedTableInfo.getSavePackageName());
+                tableInfo.setSavePackageName(finalSelectedTableInfo.getSavePackageName());
-                tableInfo.setSavePath(selectedTableInfo.getSavePath());
+                tableInfo.setSavePath(finalSelectedTableInfo.getSavePath());
             });
         }
 
@@ -244,7 +267,8 @@ public class CodeGenerateServiceImpl implements CodeGenerateService {
         Set<String> result = new TreeSet<>();
         tableInfo.getFullColumn().forEach(columnInfo -> {
             if (!columnInfo.getType().startsWith(FILTER_PACKAGE_NAME)) {
-                result.add(columnInfo.getType());
+                String type = NameUtils.getInstance().getClsFullNameRemoveGeneric(columnInfo.getType());
+                result.add(type);
             }
         });
         return result;

src/main/java/com/sjhy/plugin/service/impl/TableInfoServiceImpl.java

@@ -11,6 +11,11 @@ import com.intellij.openapi.options.ShowSettingsUtil;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.ui.Messages;
 import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.PsiClass;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiField;
+import com.intellij.psi.javadoc.PsiDocComment;
+import com.intellij.psi.javadoc.PsiDocToken;
 import com.intellij.util.ExceptionUtil;
 import com.intellij.util.containers.JBIterable;
 import com.sjhy.plugin.constants.MsgValue;
@@ -19,12 +24,27 @@ import com.sjhy.plugin.entity.SaveFile;
 import com.sjhy.plugin.entity.TableInfo;
 import com.sjhy.plugin.entity.TypeMapper;
 import com.sjhy.plugin.service.TableInfoService;
-import com.sjhy.plugin.tool.*;
+import com.sjhy.plugin.tool.CloneUtils;
-
+import com.sjhy.plugin.tool.CollectionUtil;
+import com.sjhy.plugin.tool.CurrGroupUtils;
+import com.sjhy.plugin.tool.FileUtils;
+import com.sjhy.plugin.tool.MessageDialogUtils;
+import com.sjhy.plugin.tool.NameUtils;
+import com.sjhy.plugin.tool.ProjectUtils;
+import com.sjhy.plugin.tool.StringUtils;
 import java.io.IOException;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Objects;
+import java.util.Set;
 import java.util.regex.Pattern;
 import java.util.regex.PatternSyntaxException;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * @author makejava
@@ -32,6 +52,7 @@ import java.util.regex.PatternSyntaxException;
  * @since 2018/09/02 12:13
  */
 public class TableInfoServiceImpl implements TableInfoService {
+
     /**
      * 项目对象
      */
@@ -119,21 +140,23 @@ public class TableInfoServiceImpl implements TableInfoService {
     public TableInfo getTableInfoAndConfig(DbTable dbTable) {
         TableInfo tableInfo = this.getTableInfoByDbTable(dbTable);
         // 加载配置
-        this.loadConfig(tableInfo);
+        String configFileName = getConfigFileName(tableInfo);
+        this.loadConfig(tableInfo, configFileName);
         return tableInfo;
     }
 
     /**
      * 加载单个表信息配置(用户自定义列与扩张选项)
      *
-     * @param tableInfo 表信息对象
+     * @param tableInfo      表信息对象
+     * @param configFileName 配置文件名称
      */
-    private void loadConfig(TableInfo tableInfo) {
+    private void loadConfig(TableInfo tableInfo, String configFileName) {
         if (tableInfo == null) {
             return;
         }
         // 读取配置文件中的表信息
-        TableInfo tableInfoConfig = read(tableInfo);
+        TableInfo tableInfoConfig = read(configFileName);
         // 返回空直接不处理
         if (tableInfoConfig == null) {
             return;
@@ -173,8 +196,10 @@ public class TableInfoServiceImpl implements TableInfoService {
                     if (configColumn.getType() == null) {
                         configColumn.setType(column.getType());
                     }
-                    if (!StringUtils.isEmpty(configColumn.getType())) {
+                    // 短类型
-                        // 短类型
+                    if (!StringUtils.isEmpty(column.getShortType())) {
+                        configColumn.setShortType(column.getShortType());
+                    } else if (!StringUtils.isEmpty(configColumn.getType())) {
                         configColumn.setShortType(nameUtils.getClsNameByFullName(configColumn.getType()));
                     }
                     // 表注释覆盖
@@ -187,10 +212,21 @@ public class TableInfoServiceImpl implements TableInfoService {
                     // 添加至新列表中
                     fullColumn.add(configColumn);
                     // 是否为主键
-                    if (DasUtil.isPrimary(configColumn.getObj())) {
+                    if (configColumn.getObj() != null) {
-                        pkColumn.add(configColumn);
+                        if (DasUtil.isPrimary(configColumn.getObj())) {
+                            pkColumn.add(configColumn);
+                        } else {
+                            otherColumn.add(configColumn);
+                        }
                     } else {
-                        otherColumn.add(configColumn);
+                        boolean isPk = tableInfo.getPkColumn() != null
+                                && tableInfo.getPkColumn().stream()
+                                .anyMatch(c -> Objects.equals(c.getName(), column.getName()));
+                        if (isPk) {
+                            pkColumn.add(configColumn);
+                        } else {
+                            otherColumn.add(configColumn);
+                        }
                     }
                     break;
                 }
@@ -251,12 +287,15 @@ public class TableInfoServiceImpl implements TableInfoService {
             for (TypeMapper typeMapper : typeMapperList) {
                 try {
                     // 不区分大小写查找类型
-                    if (Pattern.compile(typeMapper.getColumnType(), Pattern.CASE_INSENSITIVE).matcher(typeName).matches()) {
+                    if (Pattern.compile(typeMapper.getColumnType(), Pattern.CASE_INSENSITIVE).matcher(typeName)
+                            .matches()) {
                         continue FLAG;
                     }
                 } catch (PatternSyntaxException e) {
                     if (!errorCount.contains(typeMapper.getColumnType())) {
-                        Messages.showWarningDialog("类型映射《" + typeMapper.getColumnType() + "》存在语法错误，请及时修正。报错信息:" + e.getMessage(), MsgValue.TITLE_INFO);
+                        Messages.showWarningDialog(
+                                "类型映射《" + typeMapper.getColumnType() + "》存在语法错误，请及时修正。报错信息:" + e.getMessage(),
+                                MsgValue.TITLE_INFO);
                         errorCount.add(typeMapper.getColumnType());
                     }
                 }
@@ -280,7 +319,12 @@ public class TableInfoServiceImpl implements TableInfoService {
     @Override
     public void save(TableInfo tableInfo) {
         // 获取未修改前的原数据
-        TableInfo oldTableInfo = getTableInfoByDbTable(tableInfo.getObj());
+        TableInfo oldTableInfo;
+        if (tableInfo.getObj() != null) {
+            oldTableInfo = getTableInfoByDbTable(tableInfo.getObj());
+        } else {
+            oldTableInfo = getTableInfoByPsiClass(tableInfo.getPsiClassObj());
+        }
         // 克隆对象，防止串改，同时原始对象丢失
         tableInfo = CloneUtils.cloneByJson(tableInfo, false);
         //排除部分字段，这些字段不进行保存
@@ -340,24 +384,85 @@ public class TableInfoServiceImpl implements TableInfoService {
             return;
         }
         // 获取保存文件
-        new SaveFile(project, dir.getPath(), getConfigFileName(oldTableInfo.getObj()), content, true, false).write();
+        new SaveFile(project, dir.getPath(), getConfigFileName(oldTableInfo), content, true, false).write();
+    }
+
+    public TableInfo getTableInfoByPsiClass(PsiClass psiClass) {
+        TableInfo tableInfo = new TableInfo();
+        tableInfo.setPsiClassObj(psiClass);
+        tableInfo.setName(psiClass.getName());
+        tableInfo.setComment(parsePsiClassComment(psiClass.getDocComment()));
+        tableInfo.setFullColumn(new ArrayList<>());
+        tableInfo.setPkColumn(new ArrayList<>());
+        tableInfo.setOtherColumn(new ArrayList<>());
+        for (PsiField field : psiClass.getAllFields()) {
+            ColumnInfo columnInfo = new ColumnInfo();
+            columnInfo.setName(field.getName());
+            columnInfo.setShortType(field.getType().getPresentableText());
+            columnInfo.setType(field.getType().getCanonicalText());
+            columnInfo.setComment(parsePsiClassComment(field.getDocComment()));
+            tableInfo.getFullColumn().add(columnInfo);
+            if ("id".equals(field.getName())) {
+                tableInfo.getPkColumn().add(columnInfo);
+            } else {
+                tableInfo.getOtherColumn().add(columnInfo);
+            }
+        }
+        return tableInfo;
+    }
+
+    private String removeClassNameGeneric(String fullName) {
+        int genericIdx = fullName.indexOf('<');
+        if (genericIdx == -1) {
+            return fullName;
+        }
+        return fullName.substring(0, genericIdx);
+    }
+
+    @Override
+    public TableInfo getTableInfoAndConfigByPsiClass(PsiClass selectPsiClass) {
+        TableInfo tableInfo = this.getTableInfoByPsiClass(selectPsiClass);
+        // 加载配置
+        String configFileName = getConfigFileName(tableInfo);
+        this.loadConfig(tableInfo, configFileName);
+        return tableInfo;
+    }
+
+    @Override
+    public List<TableInfo> getTableInfoAndConfigByPsiClass(List<PsiClass> psiClassList) {
+        if (CollectionUtil.isEmpty(psiClassList)) {
+            return Collections.EMPTY_LIST;
+        }
+        List<TableInfo> tableInfoList = new ArrayList<>(psiClassList.size());
+        psiClassList.forEach(psiClass -> tableInfoList.add(this.getTableInfoAndConfigByPsiClass(psiClass)));
+        return tableInfoList;
+    }
+
+    private String parsePsiClassComment(@Nullable PsiDocComment docComment) {
+        if (docComment == null) {
+            return null;
+        }
+        return Arrays.stream(docComment.getDescriptionElements())
+                .filter(o -> o instanceof PsiDocToken)
+                .map(PsiElement::getText)
+                .findFirst()
+                .orElse(null);
     }
 
     /**
      * 读取配置文件
      *
-     * @param tableInfo 表信息对象
+     * @param fileName 配置文件名称, 不包含路径
      * @return 读取到的配置信息
      */
-    private TableInfo read(TableInfo tableInfo) {
+    private TableInfo read(String fileName) {
         // 获取保存的目录
         VirtualFile easyCodeConfigDir = getEasyCodeConfigDirectory(project);
         if (easyCodeConfigDir == null) {
             return null;
         }
         // 获取配置文件
-        String fileName = getConfigFileName(tableInfo.getObj());
+        VirtualFile configJsonFile = easyCodeConfigDir.findChild(fileName);
-        VirtualFile configJsonFile = easyCodeConfigDir.findChild(getConfigFileName(tableInfo.getObj()));
         if (configJsonFile == null) {
             return null;
         }
@@ -369,7 +474,8 @@ public class TableInfoServiceImpl implements TableInfoService {
         // 读取并解析文件
         String json = document.getText();
         if (StringUtils.isEmpty(json)) {
-            Messages.showInfoMessage(fileName + "配置文件文件为空，请尝试手动删除" + configJsonFile.getPath() + "文件！", MsgValue.TITLE_INFO);
+            Messages.showInfoMessage(fileName + "配置文件文件为空，请尝试手动删除" + configJsonFile.getPath() + "文件！",
+                    MsgValue.TITLE_INFO);
             return null;
         }
         return parser(json, configJsonFile);
@@ -419,12 +525,17 @@ public class TableInfoServiceImpl implements TableInfoService {
     /**
      * 获取配置文件名称
      *
-     * @param dbTable 表信息对象
+     * @param tableInfo 表信息对象
      * @return 对应的配置文件名称
      */
-    private String getConfigFileName(DbTable dbTable) {
+    private String getConfigFileName(TableInfo tableInfo) {
-        String schemaName = DasUtil.getSchema(dbTable);
+        if (tableInfo.getObj() != null) {
-        return schemaName + "-" + dbTable.getName() + ".json";
+            DbTable dbTable = tableInfo.getObj();
+            String schemaName = DasUtil.getSchema(dbTable);
+            return schemaName + "-" + dbTable.getName() + ".json";
+        } else {
+            return tableInfo.getName() + "_class.json";
+        }
     }
 
     /**
@@ -437,7 +548,8 @@ public class TableInfoServiceImpl implements TableInfoService {
         try {
             return objectMapper.readValue(str, TableInfo.class);
         } catch (IOException e) {
-            Messages.showWarningDialog("读取配置失败，JSON反序列化异常。请尝试手动删除" + originalFile.getPath() + "文件！", MsgValue.TITLE_INFO);
+            Messages.showWarningDialog("读取配置失败，JSON反序列化异常。请尝试手动删除" + originalFile.getPath() + "文件！",
+                    MsgValue.TITLE_INFO);
             ExceptionUtil.rethrow(e);
         }
         return null;

src/main/java/com/sjhy/plugin/tool/CacheDataUtils.java

@@ -1,6 +1,7 @@
 package com.sjhy.plugin.tool;
 
 import com.intellij.database.psi.DbTable;
+import com.intellij.psi.PsiClass;
 import lombok.Data;
 
 import java.util.List;
@@ -41,4 +42,14 @@ public class CacheDataUtils {
      * 所有选中的表
      */
     private List<DbTable> dbTableList;
+
+    /**
+     * 选中的类
+     */
+    private PsiClass selectPsiClass;
+
+    /**
+     * 所有选中的表
+     */
+    private List<PsiClass> psiClassList;
 }

src/main/java/com/sjhy/plugin/tool/NameUtils.java

@@ -89,7 +89,26 @@ public class NameUtils {
      * @return 类名
      */
     public String getClsNameByFullName(String fullName) {
-        return fullName.substring(fullName.lastIndexOf('.') + 1);
+        int genericIdx = fullName.indexOf('<');
+        if (genericIdx == -1) {
+            return fullName.substring(fullName.lastIndexOf('.') + 1);
+        }
+        String className = fullName.substring(0, genericIdx);
+        return fullName.substring(className.lastIndexOf('.') + 1);
+    }
+
+    /**
+     * 通过java全名获取类名
+     *
+     * @param fullName 全名
+     * @return 类名
+     */
+    public String getClsFullNameRemoveGeneric(String fullName) {
+        int genericIdx = fullName.indexOf('<');
+        if (genericIdx == -1) {
+            return fullName;
+        }
+        return fullName.substring(0, genericIdx);
     }
 
     /**

src/main/java/com/sjhy/plugin/ui/SelectSavePath.java

@@ -121,10 +121,23 @@ public class SelectSavePath extends JDialog {
      */
     private List<Module> moduleList;
 
+    /**
+     * 实体模式生成代码
+     */
+    private boolean entityMode;
+
     /**
      * 构造方法
      */
     public SelectSavePath(Project project) {
+        this(project, false);
+    }
+
+    /**
+     * 构造方法
+     */
+    public SelectSavePath(Project project, boolean entityMode) {
+        this.entityMode = entityMode;
         this.project = project;
         this.tableInfoService = TableInfoService.getInstance(project);
         this.codeGenerateService = CodeGenerateService.getInstance(project);
@@ -217,7 +230,12 @@ public class SelectSavePath extends JDialog {
             }
         }
         // 保存配置
-        TableInfo tableInfo = tableInfoService.getTableInfoAndConfig(cacheDataUtils.getSelectDbTable());
+        TableInfo tableInfo = null;
+        if(!entityMode) {
+            tableInfo = tableInfoService.getTableInfoAndConfig(cacheDataUtils.getSelectDbTable());
+        } else {
+            tableInfo = tableInfoService.getTableInfoAndConfigByPsiClass(cacheDataUtils.getSelectPsiClass());
+        }
         tableInfo.setSavePath(savePath);
         tableInfo.setSavePackageName(packageField.getText());
         tableInfo.setPreName(preField.getText());
@@ -229,7 +247,7 @@ public class SelectSavePath extends JDialog {
         tableInfoService.save(tableInfo);
 
         // 生成代码
-        codeGenerateService.generateByUnifiedConfig(getSelectTemplate(), unifiedConfig.isSelected(), !titleConfig.isSelected());
+        codeGenerateService.generateByUnifiedConfig(getSelectTemplate(), unifiedConfig.isSelected(), !titleConfig.isSelected(), this.entityMode);
         // 关闭窗口
         dispose();
     }
@@ -341,7 +359,13 @@ public class SelectSavePath extends JDialog {
         });
 
         // 获取选中的表信息（鼠标右键的那张表），并提示未知类型
-        TableInfo tableInfo = tableInfoService.getTableInfoAndConfig(cacheDataUtils.getSelectDbTable());
+        TableInfo tableInfo = null;
+        if(entityMode) {
+            tableInfo = tableInfoService.getTableInfoAndConfigByPsiClass(cacheDataUtils.getSelectPsiClass());
+        } else {
+            tableInfo = tableInfoService.getTableInfoAndConfig(cacheDataUtils.getSelectDbTable());
+        }
+
         // 设置默认配置信息
         if (!StringUtils.isEmpty(tableInfo.getSaveModelName())) {
             moduleComboBox.setSelectedItem(tableInfo.getSaveModelName());

src/main/resources/META-INF/plugin.xml

@@ -236,10 +236,15 @@
 
     <actions>
         <!-- Add your actions here -->
-        <group id="con.sjhy.easy.code.action" class="com.sjhy.plugin.actions.MainActionGroup" text="EasyCode" popup="true">
+        <group id="con.sjhy.easy.code.action" class="com.sjhy.plugin.actions.MainActionGroup" text="EasyCode"
+               popup="true">
             <!--生成代码菜单-->
             <add-to-group group-id="DatabaseViewPopupMenu" anchor="first"/>
         </group>
+        <action id="con.sjhy.easy.code.action.EasyCodeFromEntity"
+                class="com.sjhy.plugin.actions.EasyCodeEntityAction" text="EasyCode...">
+            <add-to-group group-id="GenerateGroup" anchor="last"/>
+        </action>
     </actions>
 
 </idea-plugin>