fix(python): a parameterized base links its class, like class Leaf(Base[str])

CHANGELOG.md

@@ -18,6 +18,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ### Fixes
 
+- In Python, a class with a parameterized base, like `class Leaf(Base[str])`, `class Repo(models.Base[int])` or `class Both(Mixin, Base[str])`, is now linked to that base, as `class Leaf(Base)` already was, and a call through a module global holding such a class, like `current = Leaf()` then `current.run()`, reaches the method the base declares. Before, such a base was skipped, so the class had no `extends` edge and anything that follows the class hierarchy missed it. Re-index Python projects after upgrading.
 - In TypeScript, an awaited call with type arguments like `await fetchApi<User>("/me")` or `await api.get<T>(url)` now links to the function it calls. Before, the call was dropped, so typed fetch wrappers looked like they had no callers. (#2586)
 - In JavaScript and TypeScript, a call or `new` in a parameter's default value, like `function renderWithNote(note = buildNote())`, a destructured `{ size = defaultSize() }` or `constructor(private svc = new Service())`, now counts as the function's own call. Before, it was dropped, so `buildNote` was missing that caller and impact on it came back incomplete. Re-index JavaScript and TypeScript projects after upgrading. Thanks @inth3shadows for the report and the pull request. (#2588)
 - In React, a component rendered at module scope, in a top-level constant or inside a test's `it(…)` callback, like `createRoot(el).render(<App />)` or `render(<Card />)`, now counts as rendered there, so an app's entry reaches its root component and a component's tests show among its callers. (#2587)

__tests__/python-generic-bases.test.ts

@@ -0,0 +1,246 @@
+/**
+ * A Python class with a parameterized base, `class Leaf(Base[str])`, extends
+ * `Base`. The grammar parses `Base[str]` in the base list as a `subscript`,
+ * and the base-list reader took only bare identifiers and attributes, so such
+ * a class had no `extends` edge at all. The reference now names the subscript's
+ * value, on the wasm path and in the kernel alike (python is kernel-routed in
+ * every published bundle, so each arm runs here; the kernel arm skips when no
+ * kernel binary is staged, like the parity suites). The resolver's own reading
+ * of a class's base list (`pythonBases`, behind module-global and
+ * `self.` dispatch) drops the type arguments the same way.
+ */
+import { describe, it, expect, beforeAll, afterAll } from 'vitest';
+import * as fs from 'fs';
+import * as os from 'os';
+import * as path from 'path';
+import { CodeGraph } from '../src';
+
+const KERNEL_PATH = path.join(
+  __dirname,
+  '..',
+  'codegraph-kernel',
+  'prebuilds',
+  `${process.platform}-${process.arch}`,
+  'codegraph-kernel.node'
+);
+
+const FILES: Record<string, string> = {
+  'app/models.py': `from typing import Generic, TypeVar
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+class Base(Generic[T]):
+    def run(self):
+        pass
+
+
+class Pair(Generic[T, U]):
+    pass
+
+
+class Mixin:
+    pass
+
+
+class Meta(type):
+    pass
+`,
+  'app/leaves.py': `from app import models
+from app.models import Base, Pair, Mixin, Meta
+
+
+class Plain(Base):
+    pass
+
+
+class Leaf(Base[str]):
+    pass
+
+
+class Nested(Base[list[int]]):
+    pass
+
+
+class Two(Pair[str, dict[str, int]]):
+    pass
+
+
+class Multi(
+    Base[
+        str
+    ]
+):
+    pass
+
+
+class Chained(Base[int][str]):
+    pass
+
+
+class Both(Mixin, Base[str]):
+    pass
+
+
+class Dotted(models.Base[int]):
+    pass
+
+
+class DottedPlain(models.Base):
+    pass
+
+
+class Keyword(Base[str], metaclass=Meta):
+    pass
+
+
+def make():
+    return Base
+
+
+class Computed(make()[0]):
+    pass
+`,
+  'app/globals.py': `from typing import Generic, TypeVar
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+class Base(Generic[T]):
+    def run(self):
+        pass
+
+
+class Pair(Generic[T, U]):
+    def run(self):
+        pass
+
+
+class Other:
+    def run(self):
+        pass
+
+
+class Plain(Base):
+    pass
+
+
+class Leaf(Base[str]):
+    pass
+
+
+class Two(Pair[str, dict[str, int]]):
+    pass
+
+
+current_a = Plain()
+current_b = Leaf()
+current_c = Two()
+
+
+def use_plain():
+    current_a.run()
+
+
+def use_leaf():
+    current_b.run()
+
+
+def use_two():
+    current_c.run()
+`,
+};
+
+const arms = [
+  { arm: 'wasm', env: { CODEGRAPH_KERNEL: '0' }, skip: false },
+  { arm: 'kernel', env: { CODEGRAPH_KERNEL_LANGS: 'python' }, skip: !fs.existsSync(KERNEL_PATH) },
+];
+
+describe.each(arms)('a parameterized Python base names its class ($arm)', ({ env, skip }) => {
+  let root = '';
+  let cg: CodeGraph;
+  const saved: Record<string, string | undefined> = {};
+
+  beforeAll(async () => {
+    if (skip) return;
+    for (const [k, v] of Object.entries(env)) {
+      saved[k] = process.env[k];
+      process.env[k] = v;
+    }
+    root = fs.mkdtempSync(path.join(os.tmpdir(), 'cg-py-generic-bases-'));
+    for (const [rel, content] of Object.entries(FILES)) {
+      fs.mkdirSync(path.dirname(path.join(root, rel)), { recursive: true });
+      fs.writeFileSync(path.join(root, rel), content);
+    }
+    fs.writeFileSync(path.join(root, 'app/__init__.py'), '');
+    cg = await CodeGraph.init(root, { index: true });
+  });
+
+  afterAll(() => {
+    cg?.close();
+    if (root) fs.rmSync(root, { recursive: true, force: true });
+    for (const [k, v] of Object.entries(saved)) {
+      if (v === undefined) delete process.env[k];
+      else process.env[k] = v;
+    }
+  });
+
+  /** Names of the classes `name` in app/leaves.py extends. */
+  const bases = (name: string): string[] => {
+    const node = cg.getNodesInFile('app/leaves.py').find((n) => n.name === name && n.kind === 'class')!;
+    expect(node, name).toBeDefined();
+    return cg.getOutgoingEdges(node.id).filter((e) => e.kind === 'extends')
+      .map((e) => cg.getNode(e.target)!.name).sort();
+  };
+
+  it.skipIf(skip)('links Base[str], Base[list[int]] and Pair[str, dict[str, int]] like a plain base', () => {
+    expect(bases('Plain')).toEqual(['Base']);
+    expect(bases('Leaf')).toEqual(['Base']);
+    expect(bases('Nested')).toEqual(['Base']);
+    expect(bases('Two')).toEqual(['Pair']);
+  });
+
+  it.skipIf(skip)('links a parameterized base written over several lines', () => {
+    expect(bases('Multi')).toEqual(['Base']);
+  });
+
+  it.skipIf(skip)('links nothing for a subscript of a subscript, Base[int][str]', () => {
+    expect(bases('Chained')).toEqual([]);
+  });
+
+  it.skipIf(skip)('resolves a method a module global inherits through a parameterized base', () => {
+    // Base, Pair and Other each declare `run`, so only the class's own bases pick
+    // one; the globals are not named after their classes, so no name match does.
+    /** Qualified names of the methods `fn` in app/globals.py calls. */
+    const calls = (fn: string): string[] => {
+      const node = cg.getNodesInFile('app/globals.py').find((n) => n.name === fn && n.kind === 'function')!;
+      expect(node, fn).toBeDefined();
+      return cg.getOutgoingEdges(node.id).filter((e) => e.kind === 'calls')
+        .map((e) => cg.getNode(e.target)!).filter((n) => n.kind === 'method').map((n) => n.qualifiedName);
+    };
+    expect(calls('use_plain')).toEqual(['Base::run']);
+    expect(calls('use_leaf')).toEqual(['Base::run']);
+    expect(calls('use_two')).toEqual(['Pair::run']);
+  });
+
+  it.skipIf(skip)('keeps every base of a list that mixes plain and parameterized ones', () => {
+    expect(bases('Both')).toEqual(['Base', 'Mixin']);
+  });
+
+  it.skipIf(skip)('links a module-qualified parameterized base like its plain form', () => {
+    expect(bases('DottedPlain')).toEqual(['Base']);
+    expect(bases('Dotted')).toEqual(['Base']);
+  });
+
+  it.skipIf(skip)('ignores a keyword argument and a base that is not a name', () => {
+    expect(bases('Keyword')).toEqual(['Base']);
+    expect(bases('Computed')).toEqual([]);
+  });
+
+  it.skipIf(skip)('links nothing for Generic[T], a class outside the project', () => {
+    const base = cg.getNodesInFile('app/models.py').find((n) => n.name === 'Base' && n.kind === 'class')!;
+    expect(cg.getOutgoingEdges(base.id).filter((e) => e.kind === 'extends')).toEqual([]);
+  });
+});

codegraph-kernel/src/python.rs

@@ -540,12 +540,20 @@ impl<'t> Walker<'t> {
         let Some(row) = self.create_node("class", &name, node, extra) else { return };
 
         // Inheritance: `class Flask(Scaffold, Mixin):` — argument_list children.
+        // A parameterized base, `class Leaf(Base[str])`, is a `subscript`
+        // whose value names the class.
         let extends_kind = edge_kind_index("extends").unwrap();
         for i in 0..node.named_child_count() {
             let Some(child) = node.named_child(i) else { continue };
             if child.kind() == "argument_list" {
                 for j in 0..child.named_child_count() {
-                    let Some(arg) = child.named_child(j) else { continue };
+                    let Some(base) = child.named_child(j) else { continue };
+                    let arg = if base.kind() == "subscript" {
+                        let Some(value) = base.child_by_field_name("value") else { continue };
+                        value
+                    } else {
+                        base
+                    };
                     if matches!(arg.kind(), "identifier" | "attribute") {
                         let name = self.text(arg).to_string();
                         self.push_ref_at(row, &name, extends_kind, arg);

src/extraction/tree-sitter.ts

@@ -7926,10 +7926,13 @@ export class TreeSitterExtractor {
       }
 
       // Python superclass list: `class Flask(Scaffold, Mixin):`
-      // argument_list contains identifier children for each parent class
+      // argument_list contains identifier children for each parent class.
+      // A parameterized base, `class Leaf(Base[str])`, is a `subscript` whose
+      // value names the class; its full text (`Base[str]`) matched nothing.
       if (child.type === 'argument_list' && node.type === 'class_definition') {
-        for (const arg of child.namedChildren) {
-          if (arg.type === 'identifier' || arg.type === 'attribute') {
+        for (const base of child.namedChildren) {
+          const arg = base.type === 'subscript' ? base.childForFieldName('value') : base;
+          if (arg && (arg.type === 'identifier' || arg.type === 'attribute')) {
             const name = getNodeText(arg, this.source);
             this.unresolvedReferences.push({
               fromNodeId: classId,

src/resolution/name-matcher.ts

@@ -398,7 +398,9 @@ function pythonRefClass(name: string, ref: UnresolvedRef, context: ResolutionCon
 function pythonBases(cls: Node, ref: UnresolvedRef, context: ResolutionContext): Node[] {
   const line = context.getFileLines?.(cls.filePath)?.[cls.startLine - 1]
     ?? context.readFile(cls.filePath)?.split('\n')[cls.startLine - 1] ?? '';
-  const bases = line.match(/^\s*class\s+\w+\s*\(([^)]*)\)/)?.[1];
+  let bases = line.match(/^\s*class\s+\w+\s*\(([^)]*)\)/)?.[1];
+  // A parameterized base, `Base[str]` or `Pair[str, Dict[str, int]]`, names its class before the brackets.
+  while (bases && /\[[^[\]]*\]/.test(bases)) bases = bases.replace(/\[[^[\]]*\]/g, '');
   return (bases?.split(',') ?? []).flatMap(name => {
     const base = pythonRefClass(name.trim(), { ...ref, filePath: cls.filePath }, context);
     return base ? [base] : [];