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

vscode/src/codeReader/rules/pydantic.test.ts

@@ -1,5 +1,5 @@
 import type { BlockView } from '../blockView'
-import type { PropertyRow, Segment } from '../contract'
+import type { IndexEntry, PropertyRow, RuleContext, Segment } from '../contract'
 import { describe, expect, it } from 'vitest'
 import { fileAttributes } from '../attributes'
 import { toBlockViews } from '../blockView'
@@ -11,7 +11,7 @@ import { celeryRule } from './celery'
 import { dataclassRule } from './dataclass'
 import { enumRule } from './enum'
 import { fastapiRule } from './fastapi'
-import { pydanticRule } from './pydantic'
+import { MODEL_INDEX, pydanticRule } from './pydantic'
 import { sqlalchemyRule } from './sqlalchemy'
 
 const py = createPythonAdapter(WASM_DIR)
@@ -20,9 +20,13 @@ const SHARE = readFixture('createShareRequest.py')
 /** 和扩展里注册的顺序一样。 */
 const RULES = [sqlalchemyRule, celeryRule, enumRule, dataclassRule, pydanticRule, fastapiRule]
 
-async function views(source: string, reports: string[] = []): Promise<BlockView[]> {
+async function views(source: string, reports: string[] = [], lookup?: RuleContext['lookup']): Promise<BlockView[]> {
   return withBlocks(py, 'python', source, blocks =>
-    toBlockViews(blocks, { source, language: py, markers: [] }, pythonZhTranslator, { rules: RULES, report: message => reports.push(message) }))
+    toBlockViews(blocks, { source, language: py, markers: [] }, pythonZhTranslator, { rules: RULES, report: message => reports.push(message), lookup }))
+}
+
+async function indexEntries(source: string): Promise<IndexEntry[]> {
+  return withBlocks(py, 'python', source, blocks => MODEL_INDEX.entries(blocks[0].node.tree.rootNode, { path: 'sample.py', source, exists: () => false }))
 }
 
 function properties(view: BlockView | undefined): PropertyRow[] {
@@ -204,6 +208,37 @@ describe('pydantic 库规则（Python）', () => {
     expect(linkedSegments(properties(share)).map(segment => segment.text)).toEqual(expect.arrayContaining(['ShareTargetType']))
   })
 
+  it('别的文件里的父类：索引认本地能确定的模型；查不到先当普通类，查到再换成数据模型卡片', async () => {
+    expect(MODEL_INDEX.mentions(AIGC)).toBe(true)
+    expect(MODEL_INDEX.mentions('from luffy_agent.model.aigc import AIGCMetadata\n\nclass Child(AIGCMetadata):\n    x: int\n')).toBe(false)
+    const entries = await indexEntries(`${AIGC}\n\nclass Plain:\n    x: str\n`)
+    expect(entries.map(entry => [entry.key, entry.value])).toEqual([['AIGCMetadata', 'model'], ['TextTrace', 'model']])
+    const alias = await indexEntries('from pydantic import BaseModel as Model\n\nclass Alias(Model):\n    x: int\n')
+    expect(alias.map(entry => entry.key)).toEqual(['Alias'])
+
+    const child = [
+      'from luffy_agent.model.aigc import AIGCMetadata',
+      '',
+      'class Child(AIGCMetadata):',
+      '    extra: int = 1',
+      '',
+      'class Other(NotAModel):',
+      '    x: int',
+      '',
+    ].join('\n')
+    const pending = await views(child)
+    expect(pending.map(view => view.kind)).toEqual(['import', 'class', 'class'])
+
+    const model = (index: { id: string }, node: { text: string }): boolean => index.id === MODEL_INDEX.id && node.text === 'AIGCMetadata'
+    const found = await views(child, [], (index, node) => model(index, node)
+      ? { key: node.text, name: { startIndex: node.startIndex, endIndex: node.endIndex }, value: 'model' }
+      : undefined)
+    const card = named(found, 'Child')
+    expect(card.kind).toBe('pydantic.model')
+    expect(propertyLines(properties(card)).slice(-2)).toEqual(['字段: 1 个', '  extra | 整数 | 1 | '])
+    expect(named(found, 'Other').kind).toBe('class')
+  })
+
   it('给模型加一个字段、改一个字段的默认值：字段行标新增，默认值那一格标改过', async () => {
     const after = AIGC
       .replace('    ReservedCode2: str\n', '    ReservedCode2: str\n    Extra: str\n')

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