Skip to content

Commit 318bd71

Browse files
committed
PyAst(fix[imports]): Capture module names for import statements
why: Plain `import x` nodes were losing module names, breaking scans and snapshots. what: - Normalize module resolution for ast.Import nodes - Narrow walkPyNodes to qualname-bearing nodes - Refresh py-ast and api-model snapshots
1 parent 626dabf commit 318bd71

4 files changed

Lines changed: 680 additions & 13 deletions

File tree

Lines changed: 314 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,314 @@
1+
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
2+
3+
exports[`buildApiPackage > builds a package model suitable for rendering 1`] = `
4+
{
5+
"generatedAt": "2025-01-01T00:00:00.000Z",
6+
"modules": [
7+
{
8+
"classes": [
9+
{
10+
"attributes": [
11+
{
12+
"annotation": "str",
13+
"docstring": null,
14+
"isPrivate": false,
15+
"kind": "variable",
16+
"location": {
17+
"colOffset": 4,
18+
"endColOffset": 13,
19+
"endLineno": 35,
20+
"lineno": 35,
21+
},
22+
"name": "name",
23+
"qualname": "sample_module.Widget.name",
24+
"summary": null,
25+
"value": null,
26+
},
27+
{
28+
"annotation": "int",
29+
"docstring": null,
30+
"isPrivate": false,
31+
"kind": "variable",
32+
"location": {
33+
"colOffset": 4,
34+
"endColOffset": 17,
35+
"endLineno": 36,
36+
"lineno": 36,
37+
},
38+
"name": "size",
39+
"qualname": "sample_module.Widget.size",
40+
"summary": null,
41+
"value": "1",
42+
},
43+
],
44+
"bases": [],
45+
"decorators": [
46+
"dataclasses.dataclass",
47+
],
48+
"docstring": "Widget model.",
49+
"isPrivate": false,
50+
"kind": "class",
51+
"location": {
52+
"colOffset": 0,
53+
"endColOffset": 41,
54+
"endLineno": 46,
55+
"lineno": 32,
56+
},
57+
"methods": [
58+
{
59+
"className": "Widget",
60+
"decorators": [],
61+
"docstring": "Return a label.
62+
63+
Returns
64+
-------
65+
str
66+
Label string.",
67+
"isAsync": false,
68+
"isPrivate": false,
69+
"kind": "method",
70+
"location": {
71+
"colOffset": 4,
72+
"endColOffset": 41,
73+
"endLineno": 46,
74+
"lineno": 38,
75+
},
76+
"name": "label",
77+
"parameters": [
78+
{
79+
"annotation": null,
80+
"default": null,
81+
"kind": "positional-or-keyword",
82+
"name": "self",
83+
"signature": "self",
84+
},
85+
],
86+
"qualname": "sample_module.Widget.label",
87+
"returns": "str",
88+
"signature": "(self)",
89+
"summary": "Return a label.",
90+
},
91+
],
92+
"name": "Widget",
93+
"qualname": "sample_module.Widget",
94+
"summary": "Widget model.",
95+
},
96+
{
97+
"attributes": [
98+
{
99+
"annotation": "int",
100+
"docstring": null,
101+
"isPrivate": false,
102+
"kind": "variable",
103+
"location": {
104+
"colOffset": 4,
105+
"endColOffset": 18,
106+
"endLineno": 52,
107+
"lineno": 52,
108+
},
109+
"name": "count",
110+
"qualname": "sample_module.Container.count",
111+
"summary": null,
112+
"value": "0",
113+
},
114+
],
115+
"bases": [],
116+
"decorators": [],
117+
"docstring": "Container type.",
118+
"isPrivate": false,
119+
"kind": "class",
120+
"location": {
121+
"colOffset": 0,
122+
"endColOffset": 19,
123+
"endLineno": 77,
124+
"lineno": 49,
125+
},
126+
"methods": [
127+
{
128+
"className": "Container",
129+
"decorators": [
130+
"property",
131+
],
132+
"docstring": "Return count of items.
133+
134+
Returns
135+
-------
136+
int
137+
Item count.",
138+
"isAsync": false,
139+
"isPrivate": false,
140+
"kind": "method",
141+
"location": {
142+
"colOffset": 4,
143+
"endColOffset": 30,
144+
"endLineno": 73,
145+
"lineno": 65,
146+
},
147+
"name": "item_count",
148+
"parameters": [
149+
{
150+
"annotation": null,
151+
"default": null,
152+
"kind": "positional-or-keyword",
153+
"name": "self",
154+
"signature": "self",
155+
},
156+
],
157+
"qualname": "sample_module.Container.item_count",
158+
"returns": "int",
159+
"signature": "(self)",
160+
"summary": "Return count of items.",
161+
},
162+
{
163+
"className": "Container",
164+
"decorators": [],
165+
"docstring": "Refresh the container.",
166+
"isAsync": true,
167+
"isPrivate": false,
168+
"kind": "method",
169+
"location": {
170+
"colOffset": 4,
171+
"endColOffset": 19,
172+
"endLineno": 77,
173+
"lineno": 75,
174+
},
175+
"name": "refresh",
176+
"parameters": [
177+
{
178+
"annotation": null,
179+
"default": null,
180+
"kind": "positional-or-keyword",
181+
"name": "self",
182+
"signature": "self",
183+
},
184+
],
185+
"qualname": "sample_module.Container.refresh",
186+
"returns": "None",
187+
"signature": "(self)",
188+
"summary": "Refresh the container.",
189+
},
190+
],
191+
"name": "Container",
192+
"qualname": "sample_module.Container",
193+
"summary": "Container type.",
194+
},
195+
],
196+
"docstring": "Sample module for AST scanning.",
197+
"exports": [],
198+
"functions": [
199+
{
200+
"decorators": [],
201+
"docstring": "Return a friendly greeting.
202+
203+
Parameters
204+
----------
205+
name : str
206+
Name to greet.
207+
208+
Returns
209+
-------
210+
str
211+
Greeting string.",
212+
"isAsync": false,
213+
"isPrivate": false,
214+
"kind": "function",
215+
"location": {
216+
"colOffset": 0,
217+
"endColOffset": 26,
218+
"endLineno": 24,
219+
"lineno": 11,
220+
},
221+
"name": "greet",
222+
"parameters": [
223+
{
224+
"annotation": "str",
225+
"default": null,
226+
"kind": "positional-or-keyword",
227+
"name": "name",
228+
"signature": "name: str",
229+
},
230+
],
231+
"qualname": "sample_module.greet",
232+
"returns": "str",
233+
"signature": "(name: str)",
234+
"summary": "Return a friendly greeting.",
235+
},
236+
{
237+
"decorators": [],
238+
"docstring": "Return provided value.
239+
240+
Parameters
241+
----------
242+
value : typing.Any
243+
Input value.
244+
245+
Returns
246+
-------
247+
typing.Any
248+
Input value.",
249+
"isAsync": false,
250+
"isPrivate": false,
251+
"kind": "function",
252+
"location": {
253+
"colOffset": 0,
254+
"endColOffset": 16,
255+
"endLineno": 93,
256+
"lineno": 80,
257+
},
258+
"name": "uses_typing",
259+
"parameters": [
260+
{
261+
"annotation": "t.Any",
262+
"default": null,
263+
"kind": "positional-or-keyword",
264+
"name": "value",
265+
"signature": "value: t.Any",
266+
},
267+
],
268+
"qualname": "sample_module.uses_typing",
269+
"returns": "t.Any",
270+
"signature": "(value: t.Any)",
271+
"summary": "Return provided value.",
272+
},
273+
],
274+
"imports": [
275+
"from __future__ import annotations",
276+
"from dataclasses import dataclasses",
277+
"from typing import typing as t",
278+
],
279+
"isPrivate": false,
280+
"kind": "module",
281+
"location": {
282+
"colOffset": 0,
283+
"endColOffset": null,
284+
"endLineno": null,
285+
"lineno": 1,
286+
},
287+
"name": "sample_module",
288+
"path": "/home/d/work/python/libtmux/astro/packages/core/py-ast/tests/fixtures/sample_module.py",
289+
"qualname": "sample_module",
290+
"summary": "Sample module for AST scanning.",
291+
"variables": [
292+
{
293+
"annotation": "str",
294+
"docstring": null,
295+
"isPrivate": false,
296+
"kind": "variable",
297+
"location": {
298+
"colOffset": 0,
299+
"endColOffset": 23,
300+
"endLineno": 8,
301+
"lineno": 8,
302+
},
303+
"name": "CONSTANT",
304+
"qualname": "sample_module.CONSTANT",
305+
"summary": null,
306+
"value": "'value'",
307+
},
308+
],
309+
},
310+
],
311+
"name": "libtmux",
312+
"root": "/home/d/work/python/libtmux/astro/packages/core/py-ast/tests/fixtures",
313+
}
314+
`;

astro/packages/core/py-ast/python/scan.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -333,9 +333,13 @@ def parse_import(node: ast.Import | ast.ImportFrom) -> dict[str, t.Any]:
333333
else:
334334
names.append(alias.name)
335335

336+
module = getattr(node, 'module', None)
337+
if module is None and isinstance(node, ast.Import) and len(node.names) == 1:
338+
module = node.names[0].name
339+
336340
return {
337341
'kind': 'import',
338-
'module': getattr(node, 'module', None),
342+
'module': module,
339343
'names': names,
340344
'level': getattr(node, 'level', None),
341345
'location': node_location(node),

astro/packages/core/py-ast/src/index.ts

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,7 @@
1-
import { execa } from 'execa'
21
import path from 'node:path'
32
import { fileURLToPath } from 'node:url'
4-
import {
5-
PyModuleArraySchema,
6-
PyModuleSchema,
7-
PyNodeSchema,
8-
type PyImport,
9-
type PyModule,
10-
type PyNode,
11-
} from './schema'
3+
import { execa } from 'execa'
4+
import { type PyImport, type PyModule, PyModuleArraySchema, PyModuleSchema, type PyNode, PyNodeSchema } from './schema'
125

136
export type PythonCommand = [string, ...string[]]
147

@@ -108,12 +101,16 @@ export const scanPythonModule = async (
108101
return PyModuleSchema.parse(modules[0])
109102
}
110103

111-
export const collectImports = (modules: PyModule[]): PyImport[] =>
112-
modules.flatMap((module) => module.imports)
104+
export const collectImports = (modules: PyModule[]): PyImport[] => modules.flatMap((module) => module.imports)
113105

114-
export function* walkPyNodes(modules: PyModule[]): Generator<PyNode> {
106+
export type PyQualifiedNode = Exclude<PyNode, PyImport>
107+
108+
export function* walkPyNodes(modules: PyModule[]): Generator<PyQualifiedNode> {
115109
for (const module of modules) {
116110
for (const item of module.items) {
111+
if (item.kind === 'import') {
112+
continue
113+
}
117114
yield item
118115
if (item.kind === 'class') {
119116
for (const method of item.methods) {

0 commit comments

Comments
 (0)