perf(python): index importlib alias rebind checks instead of rescanning

packages/adapters/python/src/usage/imports.ts

@@ -388,9 +388,26 @@
     for (let s = scope; s >= 0; s = parent[s]!) if (s === of) return true;
     return false;
   };
+  // Identifier -> statements that mention it. A statement can only rebind a
+  // name it contains, so each alias checks just those statements instead of
+  // rescanning the whole file once per alias (quadratic on many aliases).
+  const mentions = new Map<string, number[]>();
+  const identifier = /[\p{XID_Start}_][\p{XID_Continue}]*/gu;
+  texts.forEach((raw, index) => {
+    const seenHere = new Set<string>();
+    for (const form of new Set([raw, raw.normalize("NFKC")])) {
+      for (const m of form.matchAll(identifier)) seenHere.add(m[0]);
+    }
+    for (const word of seenHere) {
+      const list = mentions.get(word);
+      if (list) list.push(index);
+      else mentions.set(word, [index]);
+    }
+  });
   for (const [scope, names] of bound) {
     for (const name of [...names.keys()]) {
-      const dropped = texts.some((raw, index) => {
+      const dropped = (mentions.get(name) ?? []).some((index) => {
+        const raw = texts[index]!;
         if (!inside(scopeOf[index]!, scope)) return false;
         const own = defs.get(index)?.filter((d) => d.name === name) ?? [];
         const parsed = statementParts(raw);
@@ -449,122 +466,122 @@
 
 const clean = (dotted: string) => dotted.replace(/\s+/g, "");
 
 export const extractPythonImports: PythonImportExtractor = (source) => {
   const { statements, strings, byteStrings } = splitPythonStatements(source);
   const imports: PythonImport[] = [];
   const code: string[] = [];
   let typeCheckingIndent: number | undefined;
   let inlineSuite: { line: number; typeOnly: boolean } | undefined;
   const typeNames: TypeCheckingNames = { flags: new Set(), modules: new Set() };
   const aliasesAt = importlibAliases(statements);
 
   for (const [statementIndex, stmt] of statements.entries()) {
     if (typeCheckingIndent !== undefined && stmt.indent <= typeCheckingIndent) {
       typeCheckingIndent = undefined;
     }
     if (isTypeCheckingHeader(stmt.text, typeNames, stmt.indent === 0)) {
       // Keep the outermost guard: ending an inner block does not end it.
       typeCheckingIndent ??= stmt.indent;
       continue;
     }
     if (inlineSuite?.line !== stmt.line) inlineSuite = undefined;
     let typeOnly = typeCheckingIndent !== undefined || (inlineSuite?.typeOnly ?? false);
     let text = stmt.text;
     let conditional = stmt.indent > 0 || inlineSuite !== undefined;
     const header = clauseHeader(text);
     if (header && header.length < text.length) {
       if (isTypeCheckingHeader(header, typeNames, stmt.indent === 0)) typeOnly = true;
       text = text.slice(header.length).trimStart();
       conditional = true;
       inlineSuite = { line: stmt.endLine, typeOnly };
     }
     updateTypeCheckingNames(typeNames, stmt.text, text, !conditional);
     const base = { line: stmt.line, endLine: stmt.endLine, conditional, typeOnly };
 
     const plain = /^import\s+(.+)$/s.exec(text);
     if (plain) {
       for (const part of plain[1]!.split(",")) {
         const m = /^\s*(.+?)(?:\s+as\s+([\p{XID_Start}_][\p{XID_Continue}]*))?\s*$/su.exec(part);
         if (!m || !DOTTED.test(m[1]!.trim())) continue;
         const module = clean(m[1]!).normalize("NFKC");
         // `import a.b` binds `a`; `import a.b as c` binds `c`.
         const local = (m[2] ?? module.split(".")[0]!).normalize("NFKC");
         imports.push({ module, form: "static", names: [], local, ...base });
       }
       continue;
     }
     const from =
       /^from\s+(\.*)\s*([\p{XID_Start}_][\p{XID_Continue}.\s]*?)?\s+import\s+(.+)$/su.exec(text);
     if (from) {
       if (from[1] !== "" || from[2] === undefined || !DOTTED.test(from[2].trim())) continue;
       const names = from[3]!
         .replace(/[()]/g, " ")
         .split(",")
         .map((part) =>
           part
             .trim()
             .split(/\s+as\s+/)[0]!
             .trim(),
         )
         .filter((name) => name === "*" || IDENT.test(name))
         .map((name) => name.normalize("NFKC"));
       imports.push({ module: clean(from[2]).normalize("NFKC"), form: "static", names, ...base });
       continue;
     }
     code.push(text);
     const aliases = aliasesAt(statementIndex);
     const calls: { index: number; start: number; dunder: boolean }[] = [];
     for (const m of text.matchAll(DYNAMIC)) {
       const callee = m[0].slice(0, m[0].indexOf("(")).trim();
       calls.push({
         index: Number(m[1]),
         start: m.index + m[0].indexOf("("),
         dunder: /__import__$/.test(callee),
       });
     }
     for (const m of text.matchAll(aliases.modules)) {
       calls.push({ index: Number(m[2]), start: m.index + m[0].indexOf("("), dunder: false });
     }
     for (const m of text.matchAll(aliases.functions)) {
       const target = aliases.targets.get(m[1]!.normalize("NFKC"));
       calls.push({
         index: Number(m[2]),
         start: m.index + m[0].indexOf("("),
         dunder: target === "__import__",
       });
     }
     for (const call of calls) {
       if (byteStrings.has(call.index)) continue;
       const literal = strings[call.index] ?? "";
       if (literal === "" || literal.startsWith(".") || /\s/u.test(literal) || !DOTTED.test(literal))
         continue;
       // __import__(name, globals, locals, fromlist, level): a positive level is
       // package-relative, never an external dependency. An unknown level keeps
       // the credit, as before.
       if (call.dunder) {
         const args = callArguments(text, call.start);
         const level = args && importLevel(args);
         if (level !== undefined && level > 0) continue;
       }
       imports.push({ module: literal, form: "dynamic", names: [], ...base });
     }
   }
 
   const locals = new Set(
     imports.map((imp) => imp.local).filter((l): l is string => l !== undefined),
   );
   const attributes = new Map<string, Set<string>>();
   if (locals.size > 0) {
     const body = code.join("\n");
     for (const m of body.matchAll(
       /(?<![\p{XID_Continue}.])([\p{XID_Start}_][\p{XID_Continue}]*)\s*\.\s*([\p{XID_Start}_][\p{XID_Continue}]*)/gu,
     )) {
       const local = m[1]!.normalize("NFKC");
       if (!locals.has(local)) continue;
       let set = attributes.get(local);
       if (!set) attributes.set(local, (set = new Set()));
       set.add(m[2]!.normalize("NFKC"));
     }
   }
   return { imports, attributes };
 };

packages/adapters/python/src/usage/usage.test.ts

@@ -124,6 +124,23 @@ describe("extractPythonImports", () => {
     );
   });
 
+  it("scans many importlib aliases in linear time", () => {
+    const count = 1000;
+    const source =
+      Array.from({ length: count }, (_, i) => `import importlib as il${i}`).join("\n") +
+      '\nil0.import_module("yaml")\nil999.import_module("toml")\n' +
+      'il5 = other\nil5.import_module("dropped")\n';
+    const started = performance.now();
+    const dynamic = extractPythonImports(source)
+      .imports.filter((i) => i.form === "dynamic")
+      .map((i) => i.module);
+    const elapsed = performance.now() - started;
+    // Correctness first: a rebound alias is still dropped, the others credited.
+    assert.deepEqual(dynamic, ["yaml", "toml"]);
+    // 1000 aliases took ~5s when every alias rescanned every statement.
+    assert.ok(elapsed < 1500, `took ${Math.round(elapsed)}ms`);
+  });
+
   it("credits aliased importlib, import_module and __import__ calls", () => {
     const dyn = (src: string) =>
       extractPythonImports(src)