✨ feat(vscode): 代码阅读跨文件继承的 Pydantic 模型等索引查到后再换成数据模型卡片

vscode/src/codeReader/rules/pydantic.ts

@@ -1,4 +1,4 @@
-import type { Access, ClassBlock, CodeBlock, ImportFacts, Inline, KindIconName, LibraryRule, Node, PropertyItem, PropertyRow, RuleCard, RuleContext, ThemeColor } from '../contract'
+import type { Access, ClassBlock, CodeBlock, ImportFacts, IndexEntry, Inline, KindIconName, LibraryRule, Node, PropertyItem, PropertyRow, RuleCard, RuleContext, ThemeColor, WorkspaceIndex } from '../contract'
 import type { Call, LibraryNames } from './pythonTree'
 import { fromLibrary } from '../libraryRules'
 import { annotatedAssignments, memberDocs, withLeadingComments } from './children'
@@ -8,6 +8,7 @@ import { hint, joined, name, plain, range, word } from './segments'
 /**
  * Pydantic 的 `BaseModel` 子类整块换成「数据模型」卡片：名称 / 说明 / 访问级别 / 继承 / 选项 / 字段表，有才显示。
  * 父类链上有 pydantic 的 BaseModel（含别名、`pydantic.BaseModel`），或同一文件里继承另一个模型，都算。
+ * 父类是从别的模块导入的名字时，问工作区索引那个类是不是模型：还没查到就先当普通类，查到再换成这张卡片。
  * 字段是类主体里带注解的赋值，表和数据类同一套列；`model_config = ConfigDict(...)` 是选项，不是字段。
  * 方法、`@field_validator`、`ClassVar`、没注解的类属性照常是子块。
  */
@@ -112,15 +113,15 @@ function pydanticNames(imports: readonly ImportFacts[]): LibraryNames {
 }
 
 /** 模块顶层的类：名字和父类。装饰器包着的类也算。 */
-function moduleClasses(root: Node): { name: string, bases: Node[] }[] {
-  const classes: { name: string, bases: Node[] }[] = []
+function moduleClasses(root: Node): { name: Node, bases: Node[] }[] {
+  const classes: { name: Node, bases: Node[] }[] = []
   for (const statement of syntaxChildren(root)) {
     const definition = statement.type === 'decorated_definition' ? statement.childForFieldName('definition') : statement
     if (definition?.type !== 'class_definition')
       continue
     const className = definition.childForFieldName('name')
     if (className)
-      classes.push({ name: className.text, bases: syntaxChildren(definition.childForFieldName('superclasses')) })
+      classes.push({ name: className, bases: syntaxChildren(definition.childForFieldName('superclasses')) })
   }
   return classes
 }
@@ -129,17 +130,17 @@ function moduleClasses(root: Node): { name: string, bases: Node[] }[] {
  * 这一文件里能确定是模型的类名：父类是 pydantic 的 BaseModel，或是本文件里另一个模型。
  * 父类是从别的模块导入的名字时这里不算，那要等工作区索引。
  */
-function localModels(classes: readonly { name: string, bases: Node[] }[], names: LibraryNames): Set<string> {
+function localModels(classes: readonly { name: Node, bases: Node[] }[], names: LibraryNames): Set<string> {
   const models = new Set<string>()
   let grew = true
   while (grew) {
     grew = false
     for (const cls of classes) {
-      if (models.has(cls.name))
+      if (models.has(cls.name.text))
         continue
       const hit = cls.bases.some(base => libraryName(base, names) === 'BaseModel' || (base.type === 'identifier' && models.has(base.text)))
       if (hit) {
-        models.add(cls.name)
+        models.add(cls.name.text)
         grew = true
       }
     }
@@ -147,11 +148,70 @@ function localModels(classes: readonly { name: string, bases: Node[] }[], names:
   return models
 }
 
+/**
+ * 从语法树认 pydantic 的导入。索引建的时候没有翻译插件解析好的导入表，只能自己看导入语句。
+ * `import pydantic`、`from pydantic import BaseModel as BM`、`from pydantic.xxx import ...` 都算。
+ */
+function namesInSource(root: Node): LibraryNames {
+  const names: LibraryNames = { modules: new Set(), imported: new Map() }
+  for (const statement of syntaxChildren(root)) {
+    if (statement.type === 'import_statement') {
+      for (const item of statement.childrenForFieldName('name')) {
+        const original = item.type === 'aliased_import' ? item.childForFieldName('name') : item
+        const alias = item.type === 'aliased_import' ? item.childForFieldName('alias') : null
+        if (original && fromLibrary(original.text, LIBRARY))
+          names.modules.add(alias?.text ?? original.text)
+      }
+    }
+    if (statement.type !== 'import_from_statement')
+      continue
+    const module = statement.childForFieldName('module_name')
+    if (!module || !fromLibrary(module.text, LIBRARY))
+      continue
+    for (const item of statement.childrenForFieldName('name')) {
+      const original = item.type === 'aliased_import' ? item.childForFieldName('name') : item
+      const alias = item.type === 'aliased_import' ? item.childForFieldName('alias') : null
+      if (original)
+        names.imported.set(alias?.text ?? original.text, original.text)
+    }
+  }
+  return names
+}
+
+/** 这一文件里本地就能确定的模型，记成索引项：别的文件继承这个名字时来查。 */
+function modelEntries(root: Node): IndexEntry[] {
+  const classes = moduleClasses(root)
+  const models = localModels(classes, namesInSource(root))
+  return classes.flatMap(cls => models.has(cls.name.text) ? [{ key: cls.name.text, name: range(cls.name), value: 'model' }] : [])
+}
+
+/** 哪些类是模型。只扫源码里写了 BaseModel 的文件，子类文件不在这里解析。 */
+export const MODEL_INDEX: WorkspaceIndex = {
+  id: 'pydantic.models',
+  languageId: 'python',
+  include: '**/*.py',
+  mentions: source => /\bBaseModel\b/.test(source),
+  entries: root => modelEntries(root),
+}
+
+/**
+ * 是模型就认领。父类是别的文件里的名字时问索引：查到是模型才认，没查到先不认（宿主稍后重排）。
+ */
 function isModel(block: ClassBlock, env: Env): boolean {
   const className = block.facts.name?.text
   if (!className)
     return false
-  return localModels(moduleClasses(block.node.tree.rootNode), env.names).has(className)
+  const classes = moduleClasses(block.node.tree.rootNode)
+  if (localModels(classes, env.names).has(className))
+    return true
+  const cls = classes.find(item => item.name.text === className)
+  if (!cls)
+    return false
+  return cls.bases.some((base) => {
+    if (base.type !== 'identifier' || classes.some(item => item.name.text === base.text))
+      return false
+    return env.context.lookup(MODEL_INDEX, base)?.value === 'model'
+  })
 }
 
 /** 注解最外层的名字（`ClassVar[int]` → ClassVar），导入时的别名换回原名。 */
@@ -321,7 +381,8 @@ export const pydanticRule: LibraryRule = {
   id: LIBRARY,
   label: 'Pydantic 数据模型',
   languageIds: ['python'],
-  libraries: [LIBRARY],
+  // 子类文件常常只导入父模型、不再导入 pydantic，不能靠这个导入来决定看不看
+  libraries: [],
   claim: (block, context) => {
     if (block.kind !== 'class')
       return undefined