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 <noreply@anthropic.com>

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)

__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<string, string>, fromName: string): Promise<string[]> {
+  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<typeof n> => !!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([]);
+  });
+});
+

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;
 

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<string> {
   return pythonNameScan(context, `imports\0${filePath}\0${name}`, () => scanPythonImportKeys(name, filePath, context));