From 23ff26be68f625d25be21da3da57e75f877d8950 Mon Sep 17 00:00:00 2001 From: Eric Minish Date: Tue, 20 Oct 2026 13:04:29 -0400 Subject: [PATCH] fix(python): an aliased module import inside a function links its calls (#2607) `import a.b as X` written inside a function binds X there, ahead of any module-level binding of X, so X.f() reaches a/b.py's f, as the same import at module level already did. An undotted alias is left alone, since an optional dependency's name can match any project file. mealie's lifespan_fn now links init_db.main; the other repos are unchanged. Fixes issue 2600. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 1 + __tests__/python-local-import-alias.test.ts | 92 +++++++++++++++++++++ src/resolution/import-resolver.ts | 11 +++ src/resolution/name-matcher.ts | 25 ++++++ 4 files changed, 129 insertions(+) create mode 100644 __tests__/python-local-import-alias.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index de6d6464b..186092ccb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Fixes +- In Python, `import a.b as X` written inside a function now links its calls like `X.f()`, as the same import at module level already did. Thanks @Egrn for the report. (#2600) - In Python, `from app import email` no longer links the importing file to a project variable named `app`, such as a FastAPI or Flask `app = FastAPI()`. (#2584) - In Python, a call through a module-level instance, like `_svc.delete(obj)` after `_svc = SceneService()`, now links to the method its class has or inherits, even when other classes have a method of that name. (#2582) - In Python, a call on a module imported from outside your project, like `subprocess.run(…)` after `import subprocess`, no longer links to a project method that happens to share its name. Thanks @Egrn for the report. (#2599) diff --git a/__tests__/python-local-import-alias.test.ts b/__tests__/python-local-import-alias.test.ts new file mode 100644 index 000000000..11b3407c0 --- /dev/null +++ b/__tests__/python-local-import-alias.test.ts @@ -0,0 +1,92 @@ +/** + * `import a.b as X` written inside a function binds `X` there just as it does + * at module level, so `X.f()` reaches `a/b.py`'s `f` (#2600, the rest of + * #1626). The `from` forms inside a function already did. + */ + +import { describe, it, expect, afterEach } from 'vitest'; +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; +import CodeGraph from '../src/index'; + +let tempDir: string; +let cg: CodeGraph | null = null; + +async function calls(files: Record, fromName: string): Promise { + tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'codegraph-2600-')); + for (const [rel, source] of Object.entries(files)) { + fs.mkdirSync(path.dirname(path.join(tempDir, rel)), { recursive: true }); + fs.writeFileSync(path.join(tempDir, rel), source); + } + cg = await CodeGraph.init(tempDir, { index: true }); + cg.resolveReferences(); + const from = [...cg.getNodesByKind('function'), ...cg.getNodesByKind('method')].find((n) => n.name === fromName)!; + expect(from).toBeDefined(); + return cg + .getOutgoingEdges(from.id) + .filter((e) => e.kind === 'calls') + .map((e) => cg!.getNode(e.target)) + .filter((n): n is NonNullable => !!n) + .map((n) => `${n.name}@${n.filePath}`); +} + +afterEach(() => { + cg?.close(); + cg = null; + fs.rmSync(tempDir, { recursive: true, force: true }); +}); + +const PKG = { 'pkg/__init__.py': '', 'pkg/corpus.py': 'def corpus_pass(root):\n return root\n' }; + +describe('an aliased module import inside a function (#2600)', () => { + it.each([ + ['a function', 'def check():\n import pkg.corpus as CP\n return CP.corpus_pass("x")\n'], + ['a method', 'class T:\n def check(self):\n import pkg.corpus as CP\n return CP.corpus_pass("x")\n'], + ['an if block in a function', 'def check():\n if True:\n import pkg.corpus as CP\n return CP.corpus_pass("x")\n'], + ['a tab-indented function', 'def check():\n\timport pkg.corpus as CP\n\treturn CP.corpus_pass("x")\n'], + ])('links the call, in %s', async (_label, code) => { + expect(await calls({ ...PKG, 'test_local.py': code }, 'check')).toEqual(['corpus_pass@pkg/corpus.py']); + }); + + it('still links the module-level form', async () => { + const code = 'import pkg.corpus as CP\n\n\ndef check():\n return CP.corpus_pass("x")\n'; + expect(await calls({ ...PKG, 'test_module.py': code }, 'check')).toEqual(['corpus_pass@pkg/corpus.py']); + }); + + it('gives each function its own alias', async () => { + const files = { + 'x/__init__.py': '', 'x/y.py': 'def corpus_pass(r):\n return r\n', + 'z/__init__.py': '', 'z/w.py': 'def corpus_pass(r):\n return r\n', + 'm.py': 'def a():\n import x.y as m\n return m.corpus_pass(1)\n\n\ndef b():\n import z.w as m\n return m.corpus_pass(2)\n', + }; + expect(await calls(files, 'a')).toEqual(['corpus_pass@x/y.py']); + cg?.close(); cg = null; fs.rmSync(tempDir, { recursive: true, force: true }); + expect(await calls(files, 'b')).toEqual(['corpus_pass@z/w.py']); + }); + + it('does not lend a function\'s alias to a parameter of another function', async () => { + const code = 'def a():\n import pkg.corpus as CP\n return CP.corpus_pass(1)\n\n\ndef b(CP):\n return CP.corpus_pass(2)\n'; + expect(await calls({ ...PKG, 'm.py': code }, 'b')).toEqual([]); + }); + + it('keeps another function\'s module-level binding of the same name', async () => { + const files = { + ...PKG, + 'cp.py': 'class C:\n def corpus_pass(self, r):\n return r\n\n\nCP = C()\n', + 'm.py': 'from cp import CP\n\n\ndef b():\n return CP.corpus_pass(2)\n\n\ndef a():\n import pkg.corpus as CP\n return CP.corpus_pass(1)\n', + }; + expect(await calls(files, 'a')).toEqual(['corpus_pass@pkg/corpus.py']); + cg?.close(); cg = null; fs.rmSync(tempDir, { recursive: true, force: true }); + expect(await calls(files, 'b')).toEqual(['corpus_pass@cp.py']); + }); + + it('does not link an undotted alias of an optional dependency to a same-named project file', async () => { + const files = { + 'backends/__init__.py': '', 'backends/jinja2.py': 'def from_string(x):\n return x\n', + 'm.py': 'def f():\n try:\n import jinja2 as j\n except ImportError:\n return None\n return j.from_string("x")\n', + }; + expect(await calls(files, 'f')).toEqual([]); + }); +}); + diff --git a/src/resolution/import-resolver.ts b/src/resolution/import-resolver.ts index 1951975cb..a8548d7f5 100644 --- a/src/resolution/import-resolver.ts +++ b/src/resolution/import-resolver.ts @@ -26,6 +26,7 @@ import { arrayElementType, tsImportedTypeFile, reExportSource, + pythonFunctionModuleAlias, isRustAssociatedItem, } from './name-matcher'; import { JS_ARRAY_METHODS } from './js-builtins'; @@ -2299,6 +2300,16 @@ function resolvePythonModuleMember( const member = ref.referenceName.substring(dotIdx + 1).split('.')[0]; if (!member) return null; + // `import a.b as X` inside the calling function binds X there, ahead of any + // module-level binding of X, and the file's mappings don't hold it (#2600). + const local = pythonFunctionModuleAlias(receiver, ref, context); + if (local) { + const localPath = resolveImportPath(local, ref.filePath, ref.language, context) ?? + findPythonModuleFile(local, context, ref.filePath)?.filePath ?? null; + const target = localPath && localPath !== ref.filePath ? pythonModuleSymbol(localPath, member, context, 0) : null; + return target ? { original: ref, targetNodeId: target.id, confidence: 0.85, resolvedBy: 'import' } : null; + } + for (const imp of imports) { if (imp.localName !== receiver) continue; diff --git a/src/resolution/name-matcher.ts b/src/resolution/name-matcher.ts index 368421aa4..9e74ddf46 100644 --- a/src/resolution/name-matcher.ts +++ b/src/resolution/name-matcher.ts @@ -618,6 +618,31 @@ function pythonBindsLocally(name: string, ref: UnresolvedRef, context: Resolutio return false; } +/** + * The module `name` is bound to by an `import a.b as name` written in the + * function around `ref` — the innermost one that binds `name` at all (#2600). + * Import mappings are per file and read only module-level `import` lines, so + * such a function-local alias has none. Null when that function binds `name` + * any other way, binds it to several modules, or imports an undotted module + * (an optional dependency's name can match any project file's basename). + */ +export function pythonFunctionModuleAlias(name: string, ref: UnresolvedRef, context: ResolutionContext): string | null { + if (ref.language !== 'python' || !/^\w+$/.test(name)) return null; + const lines = pythonMemberLines(ref.filePath, context); + const statements = pythonStatementStarts(ref.filePath, context); + const scopes = context.getNodesInFile(ref.filePath).filter(n => + (n.kind === 'function' || n.kind === 'method') && n.startLine <= ref.line && n.endLine >= ref.line) + .sort((a, b) => b.startLine - a.startLine); + for (const scope of scopes) { + const bindings = pythonOwnLines(scope, ref.filePath, context).flatMap(i => pythonLineBindings(lines, i, statements, name)); + if (bindings.length === 0) continue; + const modules = new Set(bindings.map(b => b.kind === 'import' && b.key.endsWith(':*') ? b.key.slice(0, -2) : '')); + const [module] = modules; + return modules.size === 1 && module && module.includes('.') && !pythonBindsLocally(name, ref, context, false) ? module : null; + } + return null; +} + /** Distinct sources a file imports `name` from, at any scope (`from a import x` → `a:x`). */ function pythonImportKeys(name: string, filePath: string, context: ResolutionContext): Set { return pythonNameScan(context, `imports\0${filePath}\0${name}`, () => scanPythonImportKeys(name, filePath, context));