From a0f599e00b82afdcd816bbd4742c39e058b88c9a Mon Sep 17 00:00:00 2001 From: Colby McHenry Date: Mon, 6 Apr 2026 20:13:25 -0500 Subject: [PATCH] feat: Add Python class inheritance extraction for superclass relationships Addresses Python's class definition syntax where parent classes are specified in argument_list nodes (e.g. `class Child(Parent, Mixin):`). Extracts identifier and attribute children from argument_list as 'extends' references to properly model Python's inheritance patterns in the code graph. --- debug_python_ast.js | 26 ++++++++++++++++++++++++++ debug_python_ast2.js | 26 ++++++++++++++++++++++++++ src/extraction/tree-sitter.ts | 17 +++++++++++++++++ test_python_inheritance.js | 35 +++++++++++++++++++++++++++++++++++ 4 files changed, 104 insertions(+) create mode 100644 debug_python_ast.js create mode 100644 debug_python_ast2.js create mode 100644 test_python_inheritance.js diff --git a/debug_python_ast.js b/debug_python_ast.js new file mode 100644 index 0000000..edfff62 --- /dev/null +++ b/debug_python_ast.js @@ -0,0 +1,26 @@ +const { getParser, initGrammars, loadAllGrammars } = require('./dist/extraction/grammars'); + +(async () => { + await initGrammars(); + await loadAllGrammars(); + + const parser = getParser('python'); + + const code = `class Child(Parent): + pass`; + + const tree = parser.parse(code); + + function walk(node, depth = 0) { + const indent = ' '.repeat(depth); + const preview = node.text.substring(0, 30).replace(/\n/g, '\\n'); + console.log(`${indent}${node.type} [${node.startPosition.row}:${node.startPosition.column}] "${preview}"`); + + for (let i = 0; i < node.namedChildCount; i++) { + const child = node.namedChild(i); + if (child) walk(child, depth + 1); + } + } + + walk(tree.rootNode); +})(); diff --git a/debug_python_ast2.js b/debug_python_ast2.js new file mode 100644 index 0000000..b92d5f0 --- /dev/null +++ b/debug_python_ast2.js @@ -0,0 +1,26 @@ +const { getParser, initGrammars, loadAllGrammars } = require('./dist/extraction/grammars'); + +(async () => { + await initGrammars(); + await loadAllGrammars(); + + const parser = getParser('python'); + + const code = `class Child(Parent, Mixin, Base): + pass`; + + const tree = parser.parse(code); + + function walk(node, depth = 0) { + const indent = ' '.repeat(depth); + const preview = node.text.substring(0, 40).replace(/\n/g, '\\n'); + console.log(`${indent}${node.type} "${preview}"`); + + for (let i = 0; i < node.namedChildCount; i++) { + const child = node.namedChild(i); + if (child) walk(child, depth + 1); + } + } + + walk(tree.rootNode); +})(); diff --git a/src/extraction/tree-sitter.ts b/src/extraction/tree-sitter.ts index 98991b1..b37a34d 100644 --- a/src/extraction/tree-sitter.ts +++ b/src/extraction/tree-sitter.ts @@ -1241,6 +1241,23 @@ export class TreeSitterExtractor { } } + // Python superclass list: `class Flask(Scaffold, Mixin):` + // argument_list contains identifier children for each parent class + if (child.type === 'argument_list' && node.type === 'class_definition') { + for (const arg of child.namedChildren) { + if (arg.type === 'identifier' || arg.type === 'attribute') { + const name = getNodeText(arg, this.source); + this.unresolvedReferences.push({ + fromNodeId: classId, + referenceName: name, + referenceKind: 'extends', + line: arg.startPosition.row + 1, + column: arg.startPosition.column, + }); + } + } + } + // Go interface embedding: `type Querier interface { LabelQuerier; ... }` // constraint_elem wraps the embedded interface type identifier if (child.type === 'constraint_elem') { diff --git a/test_python_inheritance.js b/test_python_inheritance.js new file mode 100644 index 0000000..5168329 --- /dev/null +++ b/test_python_inheritance.js @@ -0,0 +1,35 @@ +const { extractFromSource } = require('./dist/extraction'); +const { initGrammars, loadAllGrammars } = require('./dist/extraction/grammars'); + +(async () => { + await initGrammars(); + await loadAllGrammars(); + + const code = ` +class Parent: + pass + +class Child(Parent): + pass + +class Multiple(Parent, Mixin): + pass +`; + + const result = extractFromSource('test.py', code); + + console.log('=== NODES ==='); + result.nodes.forEach(n => { + console.log(`${n.kind}: ${n.name} (line ${n.startLine})`); + }); + + console.log('\n=== UNRESOLVED REFERENCES ==='); + result.unresolvedReferences.forEach(r => { + console.log(`${r.referenceKind}: ${r.referenceName} (from ${r.fromNodeId})`); + }); + + console.log('\n=== EDGES ==='); + result.edges.forEach(e => { + console.log(`${e.kind}: ${e.source} -> ${e.target}`); + }); +})();