✨ feat(vscode): 代码阅读把推导式换成按读的顺序重排的遍历小框

vscode/src/codeReader/languages/python.test.ts

@@ -1,7 +1,8 @@
+import type { IterationFacts, Node } from '../contract'
 import type { Outline } from '../testSupport'
 import { describe, expect, it } from 'vitest'
-import { factsByLine, outline, readFixture, WASM_DIR, withBlocks } from '../testSupport'
-import { createPythonAdapter } from './python'
+import { factsByLine, nodesToText, outline, readFixture, runtime, WASM_DIR, withBlocks } from '../testSupport'
+import { createPythonAdapter, pythonIteration } from './python'
 
 const adapter = createPythonAdapter(WASM_DIR)
 
@@ -326,3 +327,90 @@ describe('python 类', () => {
     })
   })
 })
+
+describe('python 写在表达式里的遍历（推导式）', () => {
+  /** 源码里每个认得出的遍历（外层在前），节点换成源码。 */
+  async function iterations(source: string): Promise<unknown[]> {
+    const tree = await runtime.parse(adapter, 'python', source)
+    try {
+      const found: IterationFacts[] = []
+      const visit = (node: Node): void => {
+        const facts = pythonIteration(node)
+        if (facts)
+          found.push(facts)
+        for (const child of node.children) {
+          if (child)
+            visit(child)
+        }
+      }
+      visit(tree.rootNode)
+      return found.map(nodesToText)
+    }
+    finally {
+      tree.delete()
+    }
+  }
+
+  it('调用唯一的参数是生成器：那对括号是调用的，记在 sharedParentheses；参数不止一个时生成器自己带括号，没有', async () => {
+    expect(await iterations('"".join(secrets.choice(ALPHABET) for _ in range(length))')).toEqual([{
+      result: 'generator',
+      clauses: [{ kind: 'for', variables: '_', source: 'range(length)', async: false }],
+      value: 'secrets.choice(ALPHABET)',
+      sharedParentheses: { open: '(', close: ')' },
+    }])
+    const source = 'f(a, (x for x in y))'
+    expect(await iterations(source)).toEqual([{
+      result: 'generator',
+      clauses: [{ kind: 'for', variables: 'x', source: 'y', async: false }],
+      value: 'x',
+    }])
+  })
+
+  it('列表带条件、字典分出键和值、集合的变量是解构', async () => {
+    expect(await iterations('[u.id for u in users if u.is_active]')).toEqual([{
+      result: 'list',
+      clauses: [{ kind: 'for', variables: 'u', source: 'users', async: false }, { kind: 'if', condition: 'u.is_active' }],
+      value: 'u.id',
+    }])
+    expect(await iterations('{u.id: u for u in users}')).toEqual([{
+      result: 'dict',
+      clauses: [{ kind: 'for', variables: 'u', source: 'users', async: false }],
+      key: 'u.id',
+      value: 'u',
+    }])
+    expect(await iterations('{n for n, _ in pairs}')).toEqual([{
+      result: 'set',
+      clauses: [{ kind: 'for', variables: 'n, _', source: 'pairs', async: false }],
+      value: 'n',
+    }])
+  })
+
+  it('几节 for、if 照源码顺序；async for 标出异步', async () => {
+    expect(await iterations('[c for row in grid if row for c in row if c]')).toEqual([{
+      result: 'list',
+      clauses: [
+        { kind: 'for', variables: 'row', source: 'grid', async: false },
+        { kind: 'if', condition: 'row' },
+        { kind: 'for', variables: 'c', source: 'row', async: false },
+        { kind: 'if', condition: 'c' },
+      ],
+      value: 'c',
+    }])
+    expect(await iterations('[v async for v in stream() if v]')).toEqual([{
+      result: 'list',
+      clauses: [{ kind: 'for', variables: 'v', source: 'stream()', async: true }, { kind: 'if', condition: 'v' }],
+      value: 'v',
+    }])
+  })
+
+  it('套着的遍历各自认，外层取的值就是里层整个遍历', async () => {
+    expect(await iterations('[[y for y in row] for row in grid]')).toEqual([
+      { result: 'list', clauses: [{ kind: 'for', variables: 'row', source: 'grid', async: false }], value: '[y for y in row]' },
+      { result: 'list', clauses: [{ kind: 'for', variables: 'y', source: 'row', async: false }], value: 'y' },
+    ])
+  })
+
+  it('列表、字典、元组字面量和 for 语句不是', async () => {
+    expect(await iterations('x = [1, 2]\ny = {"a": 1}\nz = (a, b)\nfor x in y:\n    pass')).toEqual([])
+  })
+})

vscode/src/codeReader/languages/python.ts

@@ -1,4 +1,4 @@
-import type { Access, BlockShape, ClassFacts, CodeBlock, EntryFacts, FunctionFacts, ImportFacts, ImportItem, LanguageAdapter, Node, ParameterFacts, TextKind, TokenRole } from '../contract'
+import type { Access, BlockShape, ClassFacts, CodeBlock, EntryFacts, FunctionFacts, ImportFacts, ImportItem, IterationClause, IterationFacts, IterationResult, LanguageAdapter, Node, ParameterFacts, TextKind, TokenRole } from '../contract'
 import type { CommentRules } from './shared'
 import * as path from 'node:path'
 import { children, compoundShape, leafBlock, leafShape, listBlocks, named, sameNode, syntaxChildren } from './shared'
@@ -258,6 +258,54 @@ function dictionaryEntries(value: Node | null): EntryFacts[] | undefined {
   })
 }
 
+const ITERATION_RESULTS: Record<string, IterationResult> = {
+  generator_expression: 'generator',
+  list_comprehension: 'list',
+  set_comprehension: 'set',
+  dictionary_comprehension: 'dict',
+}
+
+/**
+ * 推导式、生成器表达式的结构，别的节点没有。几种读不成遍历的写法也没有，照常翻：
+ * `for x in a, b`（语法树认、Python 3 不认的多个来源），以及缺了变量、来源、条件或每一项的值（写到一半的代码）。
+ */
+export function pythonIteration(node: Node): IterationFacts | undefined {
+  const result = ITERATION_RESULTS[node.type]
+  const body = node.childForFieldName('body')
+  if (!result || !body)
+    return undefined
+  const clauses: IterationClause[] = []
+  for (const child of syntaxChildren(node)) {
+    if (child.type === 'for_in_clause') {
+      const variables = child.childForFieldName('left')
+      const sources = child.childrenForFieldName('right').filter(source => source !== null)
+      if (!variables || sources.length !== 1)
+        return undefined
+      clauses.push({ kind: 'for', variables, source: sources[0], async: child.children.some(token => token?.type === 'async') })
+    }
+    else if (child.type === 'if_clause') {
+      const [condition] = syntaxChildren(child)
+      if (!condition)
+        return undefined
+      clauses.push({ kind: 'if', condition })
+    }
+  }
+  if (clauses[0]?.kind !== 'for')
+    return undefined
+  const key = result === 'dict' ? body.childForFieldName('key') : undefined
+  const value = result === 'dict' ? body.childForFieldName('value') : body
+  if (!value || key === null)
+    return undefined
+  const facts: IterationFacts = { result, clauses, value, ...(key && { key }) }
+  // 调用唯一的参数是生成器表达式时只写一对括号（f(x for x in y)），语法树把这对括号算给生成器
+  const call = node.parent
+  const open = node.child(0)
+  const close = node.child(node.childCount - 1)
+  if (result === 'generator' && call?.type === 'call' && sameNode(call.childForFieldName('arguments'), node) && open && close)
+    facts.sharedParentheses = { open, close }
+  return facts
+}
+
 /** elif / else 展平成紧跟 if 的兄弟块。 */
 function ifBlocks(node: Node): CodeBlock[] {
   const blocks = [bodyBlock('if', node, bodyOf(node), 'other')]