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

CHANGELOG.md

@@ -39,6 +39,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - In Python, a from-import of a module whose name has no dot, like `from loader import Registry` with `loader.py` at the repository root, or `from mypackage import helper` in a `src/` layout, now links the importing file to that module's file, as `import loader` and `from pkg.loader import Registry` already did. Before, the file's dependency on the module itself was left out, so `codegraph affected` and impact could miss the files that import it. When the file also imported something else under the module's name, like `from django.core import mail` beside `from mail import …`, the module was linked to that instead. Re-index Python projects after upgrading.
 - 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 class or on a receiver named after it, like `Repo.save()` or `repo.save()` with `class Repo(Base[User])`, no longer links to a method of the base's type argument, like `User.save`. Before, every name in the base list counted as a supertype, so `User` was searched as if `Repo` inherited from it. The base itself, `Base`, is still searched. Re-index Python projects after upgrading. Thanks @JosefAschauer.
+- 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. Thanks @JosefAschauer.
 - In Python, a call of your own class's method named like a dictionary, list or file method, such as `get`, `update`, `items`, `pop` or `close`, now links to that method when the function says what the variable holds: `reg = SystemPackRegistry()` then `reg.get(…)`, `registry = load_registry()` with `def load_registry() -> SystemPackRegistry`, or a parameter written `reg: SystemPackRegistry`. Before, these calls were dropped as built-in methods unless the variable was named after the class. A variable the function also sets to something else, like `{}` or what `json.load` returns, a class imported from another package, like `queue.Queue`, and a union type still link nothing. Re-index Python projects after upgrading. Thanks @inth3shadows for the report. (#2583)
 - In Python, Scala and Java, a standard-library method called on a value whose type CodeGraph can't see, like `options.setdefault(…)` on a dictionary, `m.getOrElse(…)` on an `Option` or `data.put(…)` on a map held in a `var`, no longer links to an unrelated project method that happens to share its name; the call still links when the variable is named after your class, like `registry.get(…)` for a `Registry`. Re-index Python, Scala and Java projects after upgrading. Thanks @danusha2345. (#1947)
 - 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-generic-bases.test.ts

@@ -0,0 +1,247 @@
+/**
+ * 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 kernelAvailable = fs.existsSync(KERNEL_PATH) || process.env.CODEGRAPH_KERNEL_EXPECT === '1';
+
+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: !kernelAvailable },
+];
+
+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

@@ -8058,10 +8058,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

@@ -413,7 +413,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] : [];