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

A Python class with a parameterized base had no extends edge. The grammar
parses Base[str] in a base list as a subscript, and the base-list reader took
only bare identifiers and attributes, so the base was skipped, on the wasm path
and in the kernel alike. Everything that follows the class hierarchy missed
such a class. The resolver's own text reader of a base list, pythonBases, also
looked up Base[str] verbatim and split Pair[str, int] at its inner comma, so
current = Leaf() then current.run() did not reach Base.run.

A subscript base is now unwrapped to its value before the identifier or
attribute check, in src/extraction/tree-sitter.ts and in the kernel's
python.rs. pythonBases drops bracketed type arguments before splitting the list.
Generic[T] and Protocol[T] stay outside the project like their plain forms; a
computed base, a subscript of a subscript and keyword arguments link nothing.

Validation:
- The test file runs on both arms: red on main's source and kernel (12 of 16
  fail, 6 per arm), green here (16 passed).
- Kernel parity with this kernel on 3,155 Python files of flask, requests,
  mypy, litestar, sqlalchemy, polar and django: 0 diffs, and the same 5 files
  defer to wasm as on main. The kernel and wasm edge dumps of the repos run
  both ways are identical.
- A/B against main on 7 Python repos, 457,948 edges: 1,492 added (901 extends,
  555 calls, 36 references), 0 removed, 607 call sites moved, 223 kept their
  target with higher confidence.
  polar: 491 self.get_one_or_none() style calls move from the protocol stub
  RepositoryProtocol at 0.4 to RepositoryBase and its mixin at 0.9.
  litestar: 12 calls in GenericSyncMockRepository move from
  AbstractAsyncRepository to AbstractSyncRepository.
  sqlalchemy: 435 extends and 554 calls added, like t.insert() reaching
  TableClause.insert.
  We read 44 sampled changes and every small bucket. They read right, except
  12 self.op() sites that move to the typing-only overload SQLCoreOperations.op
  and one Enum.copy that moves to TypeEngine.copy.
- 3 of the 901 extends are wrong: sqlalchemy's Sequence[_T], Tuple[int, ...]
  and MutableSet[_T] from typing link to its own classes of those names. The
  plain form does this on main (28 class X(Enum) there), and none of the 3
  brought a call edge.
- Controls express and gin: full dumps byte-identical.
- Index cost with PR #3001 on top, sqlalchemy, median of 3 interleaved: 23.7 s
  on main, 25.2 s here, with 1,143 more edges.

Original PR #2998 by @JosefAschauer.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Josh66
2026-10-30 20:05:10 +00:00
committed by Colby Mchenry
co-authored by Claude Opus 5.5
parent 9bfdd09238
commit 3fca009e7b
5 changed files with 266 additions and 5 deletions
+1
View File
@@ -40,6 +40,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)
+247
View File
@@ -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([]);
});
});
+9 -1
View File
@@ -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);
+6 -3
View File
@@ -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,
+3 -1
View File
@@ -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] : [];