python: resolve calls through an imported module name (M.f() / from a import b as M)

src/graph/extract.ts

@@ -335,6 +335,12 @@ export interface WalkCtx {
   enclosingClass: string | null; // nearest enclosing class (py/ts `self`/`this`)
   goReceiverVar: string | null; // Go receiver var, e.g. `w` in `func (w *Worker)`
   importedSymbols: ReadonlyMap<string, { name: string; specifier: string }>;
+  /** Python only: local name → dotted module path, for every import that binds a
+   * MODULE (`import a.b as M`, `from a import b as M`, `from a import b`,
+   * `import a.b`). A call whose receiver is one of these names is a call into
+   * that module's file, not a method on an object of unknown type — see the
+   * module-receiver branch in walk() and resolvePyModule in resolve.ts. */
+  importedModules: ReadonlyMap<string, string>;
   // R6 (Phase 2): which list we're inside while walking an `R6Class(...)` call's
   // arguments — set only for the direct span of a `public =`/`private =`/
   // `active =` `list(...)`'s own entries (see walk()'s special-cased `argument`
@@ -389,6 +395,7 @@ export function extractFile(rel: string, source: string, lang: Language): Extrac
   const root = parseSource(source);
   const bindings = collectBindings(root, lang);
   const importedSymbols = collectImportedSymbols(root, lang);
+  const importedModules = lang === "python" ? collectPyImportedModules(root) : EMPTY_MODULES;
   const rGenerics = lang === "r" ? collectRGenerics(root) : EMPTY_SET;
 
   const nodes: NodeV1[] = [
@@ -422,6 +429,7 @@ export function extractFile(rel: string, source: string, lang: Language): Extrac
     enclosingClass: null,
     goReceiverVar: null,
     importedSymbols,
+    importedModules,
     rR6Access: null,
     rGenerics,
     rSuperClass: null,
@@ -766,7 +774,17 @@ function walk(node: Parser.SyntaxNode, ctx: WalkCtx, out: NodeV1[], edges: RawEd
         });
       } else {
         const recvType = resolveRecvType(callee.receiver, ctx);
-        edges.push(recvType ? { ...callEdge, recvType } : callEdge);
+        // Python: `M.f()` where `M` is an imported MODULE (not a bound type). The
+        // import names the exact file, so the call carries that module as its
+        // specifier and resolve.ts looks the name up inside that file only. A
+        // receiver that resolved to a type keeps the typed-member path — a class
+        // alias (`from pkg import Widget as W; W.make()`) is not a module.
+        const moduleSpec =
+          !recvType && ctx.lang === "python" && callee.viaMember && callee.receiver
+            ? ctx.importedModules.get(callee.receiver)
+            : undefined;
+        if (moduleSpec) edges.push({ ...callEdge, viaMember: false, specifier: moduleSpec });
+        else edges.push(recvType ? { ...callEdge, recvType } : callEdge);
       }
     }
   } else if (ctx.lang === "php" && node.type === "use_declaration") {
@@ -2397,10 +2415,74 @@ function pyReceiver(fn: Parser.SyntaxNode): string | undefined {
     const innerObj = obj.childForFieldName("object");
     const innerAttr = obj.childForFieldName("attribute");
     if (innerObj?.type === "identifier" && innerObj.text === "self" && innerAttr) return `self.${innerAttr.text}`;
+    // A pure dotted chain of names (`pkg.mod.f()`) is spelled the way
+    // `import pkg.mod` binds it; anything else on the chain (a call, a subscript)
+    // carries no local clue and stays unknown.
+    const dotted = pyDottedName(obj);
+    if (dotted && !dotted.startsWith("self.") && !dotted.startsWith("cls.")) return dotted;
   }
   return undefined;
 }
 
+/** `a.b.c` as text when the node is only identifiers joined by attribute access. */
+function pyDottedName(node: Parser.SyntaxNode): string | null {
+  if (node.type === "identifier") return node.text;
+  if (node.type !== "attribute") return null;
+  const obj = node.childForFieldName("object");
+  const attr = node.childForFieldName("attribute");
+  if (!obj || !attr) return null;
+  const head = pyDottedName(obj);
+  return head ? `${head}.${attr.text}` : null;
+}
+
+const EMPTY_MODULES: ReadonlyMap<string, string> = new Map();
+
+/**
+ * Python import statements that bind a module to a local name:
+ *   import a.b            → "a.b" → "a.b"
+ *   import a.b as M       → "M"   → "a.b"
+ *   from a import b       → "b"   → "a.b"   (b may be a symbol, not a module —
+ *                                             resolve.ts drops it when no file matches)
+ *   from a import b as M  → "M"   → "a.b"
+ *   from . import b       → "b"   → ".b"    (relative; resolved against the file's dir)
+ * `from a import *` binds nothing nameable and is skipped.
+ */
+function collectPyImportedModules(root: Parser.SyntaxNode): Map<string, string> {
+  const out = new Map<string, string>();
+  const visit = (node: Parser.SyntaxNode): void => {
+    if (node.type === "import_statement") {
+      for (const c of node.namedChildren) {
+        if (c.type === "dotted_name") out.set(c.text, c.text);
+        else if (c.type === "aliased_import") {
+          const name = c.childForFieldName("name")?.text;
+          const alias = c.childForFieldName("alias")?.text;
+          if (name && alias) out.set(alias, name);
+        }
+      }
+      return;
+    }
+    if (node.type === "import_from_statement") {
+      const moduleNode = node.childForFieldName("module_name");
+      const base = moduleNode?.text;
+      if (!base) return;
+      const join = (name: string): string => (base.endsWith(".") ? `${base}${name}` : `${base}.${name}`);
+      for (const c of node.namedChildren) {
+        if (c === moduleNode || sameSyntaxNode(c, moduleNode)) continue;
+        if (c.type === "dotted_name") out.set(c.text, join(c.text));
+        else if (c.type === "aliased_import") {
+          const name = c.childForFieldName("name")?.text;
+          const alias = c.childForFieldName("alias")?.text;
+          if (name && alias) out.set(alias, join(name));
+        }
+      }
+      return;
+    }
+    for (const child of node.namedChildren) visit(child);
+  };
+  visit(root);
+  return out;
+}
+
 /**
  * PHP call shapes: `foo()` (function_call_expression), `$obj->m()` /
  * `$obj?->m()` (member/nullsafe_member_call_expression), and `Cls::m()`

src/graph/resolve.ts

@@ -268,6 +268,20 @@ export function resolveEdges(
         if (hit && hit.id !== e.source) add(e.source, hit.id, "references", hit.confidence);
       }
     } else if (e.relation === "calls") {
+      if (e.specifier) {
+        // Python `M.f()` with `M` an imported module (extract.ts stamps the
+        // module's dotted path as the specifier). The import names the file, so
+        // the lookup is confined to it: a same-named function anywhere else in
+        // the repo is not a candidate, and a name defined twice in that file (or
+        // not at all — `from a import b` where `b` is a symbol) drops.
+        const targetFile = resolvePyModule(e.specifier, e.file, byId);
+        if (!targetFile) continue;
+        const candidates = (perFileName.get(targetFile)?.get(e.name!) ?? []).filter((n) =>
+          PY_MODULE_CALL_KINDS.includes(n.kind),
+        );
+        if (candidates.length === 1) add(e.source, candidates[0].id, "calls", "extracted");
+        continue;
+      }
       if (e.viaMember) {
         if (!e.recvType) continue;
         const hit = resolveTypedMember(e.recvType, e.name!, e.file, ownerMethod, classParents, classTraits, e.argCount);
@@ -336,6 +350,44 @@ export function resolveEdges(
   return out;
 }
 
+/** What `M.name()` may target inside a module: a function, or a class
+ * (`M.Widget()` constructs). Methods are never module-level. */
+const PY_MODULE_CALL_KINDS: Kind[] = ["function", "class"];
+
+/**
+ * Resolve a Python dotted module path to an in-repo file node id, or null.
+ *
+ * Absolute (`a.b`): tried as `a/b.py` then `a/b/__init__.py` under every
+ * ancestor directory of the importing file, nearest first — Python's own
+ * search order (the script's directory, then the roots on sys.path), which also
+ * covers a `src/` layout without parsing any build file.
+ * Relative (`.b`, `..b`): each leading dot after the first climbs one directory
+ * from the importing file's own.
+ * A path that matches nothing is an external package — null, never a guess.
+ */
+function resolvePyModule(spec: string, file: string, byId: Map<string, NodeV1>): string | null {
+  const dir = posix.dirname(toPosixPath(file));
+  const asPath = (base: string, dotted: string): string[] => {
+    const rel = dotted.split(".").filter(Boolean).join("/");
+    const stem = base === "." || base === "" ? rel : posix.join(base, rel);
+    return [`${stem}.py`, `${stem}/__init__.py`];
+  };
+  const hit = (cands: string[]): string | null => cands.find((c) => byId.has(c)) ?? null;
+  const dots = spec.match(/^\.+/)?.[0].length ?? 0;
+  if (dots > 0) {
+    let base = dir;
+    for (let i = 1; i < dots; i++) base = posix.dirname(base);
+    return hit(asPath(base, spec.slice(dots)));
+  }
+  let base = dir;
+  for (;;) {
+    const found = hit(asPath(base, spec));
+    if (found) return found;
+    if (base === "." || base === "" || base === "/") return null;
+    base = posix.dirname(base);
+  }
+}
+
 function push<T>(map: Map<string, T[]>, key: string, val: T): void {
   const arr = map.get(key);
   if (arr) arr.push(val);

test/graph-python-module-alias.test.ts

@@ -0,0 +1,129 @@
+/**
+ * Python calls through an imported MODULE name.
+ *
+ * `from analysis import pump_signature as P` followed by `P.find_pumps()` is the
+ * dominant call style in numeric Python codebases (numpy-style `import x as np`
+ * generalised to the project's own modules). The extractor saw the receiver `P`,
+ * found no type binding for it and — correctly, for an object of unknown class —
+ * dropped the edge rather than guess. The result was `graft callers find_pumps`
+ * reporting "no indexed callers" for a function called from a dozen sites.
+ *
+ * A module receiver is not an unknown object: the import statement names the
+ * exact file. So the call resolves inside THAT file only, never globally — a
+ * same-named function elsewhere in the repo must not become an edge (the last
+ * test pins that). Four import spellings bind a local module name:
+ *
+ *   import pkg.mod as M          → M
+ *   from pkg import mod as M     → M
+ *   from pkg import mod          → mod
+ *   import pkg.mod               → pkg.mod   (call spelled `pkg.mod.f()`)
+ */
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import { mkdtempSync, writeFileSync, mkdirSync, rmSync } from "node:fs";
+import { join } from "node:path";
+import { tmpdir } from "node:os";
+import { buildGraph } from "../src/graph/build.js";
+import { readGraph, wiringPath } from "../src/graph/write.js";
+import type { GraphV1 } from "../src/graph/types.js";
+
+const MOD = `def find_pumps(close):
+    return close
+
+
+def compute_all(x):
+    return x
+`;
+
+/** A package whose __init__ is the module: \`from pkg import init_mod\` must reach it. */
+const INIT_MOD = `def in_init():
+    return 1
+`;
+
+/** Same-named function in an unrelated file: must never be the target. */
+const DECOY = `def find_pumps(other):
+    return other
+`;
+
+const CALLERS = `import pkg.mod as M
+from pkg import mod as P
+from pkg import mod
+from pkg import init_mod
+import pkg.mod
+
+
+def via_import_as():
+    return M.find_pumps(1)
+
+
+def via_from_as():
+    return P.find_pumps(2)
+
+
+def via_from_plain():
+    return mod.compute_all(3)
+
+
+def via_dotted():
+    return pkg.mod.compute_all(4)
+
+
+def via_init():
+    return init_mod.in_init()
+`;
+
+/** Receiver that is NOT a module: the drop rule must still hold. */
+const UNKNOWN = `def use(obj):
+    return obj.find_pumps(5)
+`;
+
+function makeFixture(): string {
+  const dir = mkdtempSync(join(tmpdir(), "graft-pyalias-"));
+  mkdirSync(join(dir, "pkg", "init_mod"), { recursive: true });
+  writeFileSync(join(dir, "pkg", "__init__.py"), "");
+  writeFileSync(join(dir, "pkg", "mod.py"), MOD);
+  writeFileSync(join(dir, "pkg", "init_mod", "__init__.py"), INIT_MOD);
+  writeFileSync(join(dir, "decoy.py"), DECOY);
+  writeFileSync(join(dir, "callers.py"), CALLERS);
+  writeFileSync(join(dir, "unknown.py"), UNKNOWN);
+  return dir;
+}
+
+async function buildFixture(dir: string): Promise<GraphV1> {
+  await buildGraph(dir);
+  const graph = readGraph(wiringPath(join(dir, "graft")));
+  assert.ok(graph, "wiring graph should be written");
+  return graph!;
+}
+
+function hasCall(graph: GraphV1, source: string, target: string): boolean {
+  return graph.edges.some((e) => e.relation === "calls" && e.source === source && e.target === target);
+}
+
+test("Python: calls through an imported module name resolve into that module's file", async () => {
+  const dir = makeFixture();
+  try {
+    const g = await buildFixture(dir);
+    assert.ok(hasCall(g, "callers.py#via_import_as", "pkg/mod.py#find_pumps"), "import pkg.mod as M → M.f()");
+    assert.ok(hasCall(g, "callers.py#via_from_as", "pkg/mod.py#find_pumps"), "from pkg import mod as P → P.f()");
+    assert.ok(hasCall(g, "callers.py#via_from_plain", "pkg/mod.py#compute_all"), "from pkg import mod → mod.f()");
+    assert.ok(hasCall(g, "callers.py#via_dotted", "pkg/mod.py#compute_all"), "import pkg.mod → pkg.mod.f()");
+    assert.ok(hasCall(g, "callers.py#via_init", "pkg/init_mod/__init__.py#in_init"), "package module via __init__.py");
+  } finally {
+    rmSync(dir, { recursive: true, force: true });
+  }
+});
+
+test("Python: a module-qualified call never resolves to a same-named function elsewhere", async () => {
+  const dir = makeFixture();
+  try {
+    const g = await buildFixture(dir);
+    const toDecoy = g.edges.filter((e) => e.relation === "calls" && e.target === "decoy.py#find_pumps");
+    assert.deepEqual(toDecoy, [], "decoy.py#find_pumps must have no callers");
+    // Unknown receiver: still dropped, not guessed (the #35 precision rule).
+    const fromUnknown = g.edges.filter((e) => e.relation === "calls" && e.source === "unknown.py#use");
+    assert.deepEqual(fromUnknown, [], "obj.find_pumps() on an untyped receiver stays unresolved");
+  } finally {
+    rmSync(dir, { recursive: true, force: true });
+  }
+});