From 597ab945625af0f721146f6f6da39e1cab5995bf Mon Sep 17 00:00:00 2001 From: cruldra Date: Fri, 2 Oct 2026 00:12:05 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat(vscode):=20=E4=BB=A3=E7=A0=81?= =?UTF-8?q?=E9=98=85=E8=AF=BB=E8=B7=A8=E6=96=87=E4=BB=B6=E7=BB=A7=E6=89=BF?= =?UTF-8?q?=E7=9A=84=20Pydantic=20=E6=A8=A1=E5=9E=8B=E7=AD=89=E7=B4=A2?= =?UTF-8?q?=E5=BC=95=E6=9F=A5=E5=88=B0=E5=90=8E=E5=86=8D=E6=8D=A2=E6=88=90?= =?UTF-8?q?=E6=95=B0=E6=8D=AE=E6=A8=A1=E5=9E=8B=E5=8D=A1=E7=89=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- vscode/src/codeReader/rules/pydantic.test.ts | 43 ++++++++++- vscode/src/codeReader/rules/pydantic.ts | 79 +++++++++++++++++--- 2 files changed, 109 insertions(+), 13 deletions(-) diff --git a/vscode/src/codeReader/rules/pydantic.test.ts b/vscode/src/codeReader/rules/pydantic.test.ts index 5b53843..4374538 100644 --- a/vscode/src/codeReader/rules/pydantic.test.ts +++ b/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 { +async function views(source: string, reports: string[] = [], lookup?: RuleContext['lookup']): Promise { 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 { + 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') diff --git a/vscode/src/codeReader/rules/pydantic.ts b/vscode/src/codeReader/rules/pydantic.ts index 9812b84..9f9ff59 100644 --- a/vscode/src/codeReader/rules/pydantic.ts +++ b/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 { +function localModels(classes: readonly { name: Node, bases: Node[] }[], names: LibraryNames): Set { const models = new Set() 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