feat: Add Python class inheritance extraction for superclass relationships

colbymchenry · colbymchenry · commit a0f599e00b82 · 2026-04-06T20:13:25.000-05:00
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.
diff --git a/debug_python_ast.js 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
@@ -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
@@ -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
@@ -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}`);
+  });
+})();
