Support inheritance detection in Python Extractor (#175)

* test(python): cover canonical inheritance targets

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(python): resolve inheritance targets canonically

This is the resolution of the bug that inheritance in Python was not
resolved to an "implements" edge.

The fix adds a post-processing step to the Python extractor:
  + Collect all parsed symbols and index them by short class name.
  + Resolve unambiguous inheritance targets to fully qualified symbols.
  + Preserve ambiguous targets unchanged to avoid incorrect edges.

* fix(python): use imports to resolve inheritance bases

Resolving ambiguous symbol names during inheritance detection by taking
the import context into account.

* fix(python): preserve inheritance fact compatibility

Fixing regression in the golden test.

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: jnthnhrrr

internal/extractors/pythonextractor/python.go

@@ -176,6 +176,7 @@ func (e *PythonExtractor) Extract(ctx context.Context, repoPath string, files []
 		fileModules[strings.TrimSuffix(f, ".py")] = true
 	}
 	resolveCallTargets(allFacts, fileModules, pkgDirs)
+	resolveImplementsTargets(allFacts, fileModules, pkgDirs)
 
 	// Fold FastAPI include_router mount prefixes onto the bare decorator paths, so
 	// a route reads as the path it actually serves ("/api/v1/cognify") rather than

internal/extractors/pythonextractor/python_ast.go

@@ -477,6 +477,7 @@ func (w *pyWalker) handleFromImport(node *sitter.Node) {
 	// flag them as orphans. Kept only here to avoid bloating every from-import.
 	isInit := w.relFile == "__init__.py" || strings.HasSuffix(w.relFile, "/__init__.py")
 	var reexported []string
+	bindings := make(map[string]string)
 
 	// Map each imported name to a resolvable target or "" (external).
 	for i := uint(0); i < uint(node.ChildCount()); i++ {
@@ -501,6 +502,9 @@ func (w *pyWalker) handleFromImport(node *sitter.Node) {
 			importedName = pyText(c, w.src)
 			localName = importedName
 		}
+		if localName != "" && importedName != "" && importedName != "*" {
+			bindings[localName] = importedName
+		}
 
 		if isInit && importedName != "" && importedName != "*" {
 			reexported = append(reexported, importedName)
@@ -522,6 +526,9 @@ func (w *pyWalker) handleFromImport(node *sitter.Node) {
 		}
 	}
 
+	if len(bindings) > 0 {
+		depProps["bindings"] = bindings
+	}
 	if len(reexported) > 0 {
 		depProps["reexports"] = reexported
 	}

internal/extractors/pythonextractor/resolve.go

@@ -113,6 +113,86 @@ func resolveCallTargets(allFacts []facts.Fact, fileModules map[string]bool, pkgD
 	}
 }
 
+func resolveImplementsTargets(allFacts []facts.Fact, fileModules map[string]bool, pkgDirs map[string]bool) {
+	fileIdx := buildSuffixIndex(fileModules, pkgDirs)
+	topPkgs := importableRoots(fileModules, pkgDirs)
+	reexports := buildReexportIndex(allFacts, pkgDirs)
+	symbols := make(map[string]bool)
+	byShortName := make(map[string][]string)
+	importedByFile := make(map[string]map[string]string)
+	for i := range allFacts {
+		if allFacts[i].Kind == facts.KindDependency && allFacts[i].Props["from"] == true {
+			var module string
+			for _, rel := range allFacts[i].Relations {
+				if rel.Kind == facts.RelImports {
+					module = rel.Target
+					break
+				}
+			}
+			if module != "" {
+				if importedByFile[allFacts[i].File] == nil {
+					importedByFile[allFacts[i].File] = make(map[string]string)
+				}
+				for local, imported := range stringMapProp(allFacts[i].Props, "bindings") {
+					target := module + "." + imported
+					if previous, exists := importedByFile[allFacts[i].File][local]; exists && previous != target {
+						importedByFile[allFacts[i].File][local] = ""
+					} else {
+						importedByFile[allFacts[i].File][local] = target
+					}
+				}
+			}
+		}
+		if allFacts[i].Kind != facts.KindSymbol {
+			continue
+		}
+		name := allFacts[i].Name
+		symbols[name] = true
+		short := name
+		if dot := strings.LastIndexByte(name, '.'); dot >= 0 {
+			short = name[dot+1:]
+		}
+		byShortName[short] = append(byShortName[short], name)
+	}
+	for i := range allFacts {
+		if allFacts[i].Kind == facts.KindDependency && allFacts[i].Props != nil {
+			delete(allFacts[i].Props, "bindings")
+		}
+	}
+
+	for i := range allFacts {
+		f := &allFacts[i]
+		if f.Kind != facts.KindSymbol {
+			continue
+		}
+		for j := range f.Relations {
+			rel := &f.Relations[j]
+			if rel.Kind != facts.RelImplements {
+				continue
+			}
+			if symbols[rel.Target] {
+				continue
+			}
+			if strings.ContainsRune(rel.Target, '.') {
+				if resolved, keep := resolveDottedTarget(rel.Target, fileIdx, topPkgs, fileDir(f.File), reexports, symbols); keep && symbols[resolved] {
+					rel.Target = resolved
+				}
+				continue
+			}
+			if imported := importedByFile[f.File][rel.Target]; imported != "" {
+				if resolved, keep := resolveDottedTarget(imported, fileIdx, topPkgs, fileDir(f.File), reexports, symbols); keep && symbols[resolved] {
+					rel.Target = resolved
+				}
+				continue
+			}
+			candidates := byShortName[rel.Target]
+			if len(candidates) == 1 {
+				rel.Target = candidates[0]
+			}
+		}
+	}
+}
+
 // isDottedCallTarget reports whether a call target is an unresolved dotted path
 // (e.g. "a.b.c") rather than an already-resolved slash symbol name
 // ("dir/mod.sym") or a bare short name ("Foo").
@@ -242,6 +322,23 @@ func stringSliceProp(props map[string]any, key string) []string {
 	return nil
 }
 
+func stringMapProp(props map[string]any, key string) map[string]string {
+	out := make(map[string]string)
+	switch v := props[key].(type) {
+	case map[string]string:
+		for k, value := range v {
+			out[k] = value
+		}
+	case map[string]any:
+		for k, value := range v {
+			if s, ok := value.(string); ok {
+				out[k] = s
+			}
+		}
+	}
+	return out
+}
+
 // resolveDottedTarget maps a dotted call target ("a.b.c.sym") to a canonical slash
 // symbol name when its module prefix resolves to an internal file or to a package
 // that re-exports the symbol, keeps it dotted when the prefix is internal but

internal/extractors/pythonextractor/resolve_test.go

@@ -283,6 +283,83 @@ func TestResolveCallTargets_AbsoluteInternal_RewritesToSlashSymbol(t *testing.T)
 	}
 }
 
+func TestExtract_PythonResolvesInheritanceTarget(t *testing.T) {
+	dir := t.TempDir()
+	files := map[string]string{
+		"pkg/__init__.py":    "",
+		"pkg/wrongparent.py": "class Parent:\n    pass\n",
+		"pkg/parent.py":      "class Parent:\n    pass\n",
+		"pkg/child.py":       "from pkg.parent import Parent\n\nclass Child(Parent):\n    pass\n",
+	}
+	var rel []string
+	for name, content := range files {
+		full := filepath.Join(dir, name)
+		if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil {
+			t.Fatal(err)
+		}
+		if err := os.WriteFile(full, []byte(content), 0o644); err != nil {
+			t.Fatal(err)
+		}
+		rel = append(rel, name)
+	}
+
+	all, err := New().Extract(context.Background(), dir, rel)
+	if err != nil {
+		t.Fatalf("Extract: %v", err)
+	}
+
+	for _, f := range all {
+		if f.Kind != facts.KindSymbol || f.Name != "pkg/child.Child" {
+			continue
+		}
+		if hasRel(f, facts.RelImplements, "pkg/parent.Parent") {
+			return
+		}
+		t.Fatalf("Child inheritance target = %v, want pkg/parent.Parent", f.Relations)
+	}
+	t.Fatal("missing pkg/child.Child symbol")
+}
+
+func TestExtract_PythonResolvesMultipleInheritanceTargets(t *testing.T) {
+	dir := t.TempDir()
+	files := map[string]string{
+		"pkg/__init__.py": "",
+		"pkg/father.py":   "class Father:\n    pass\n",
+		"pkg/mother.py":   "class Mother:\n    pass\n",
+		"pkg/child.py":    "from pkg.father import Father\nfrom pkg.mother import Mother\n\nclass Child(Father, Mother):\n    pass\n",
+	}
+	var rel []string
+	for name, content := range files {
+		full := filepath.Join(dir, name)
+		if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil {
+			t.Fatal(err)
+		}
+		if err := os.WriteFile(full, []byte(content), 0o644); err != nil {
+			t.Fatal(err)
+		}
+		rel = append(rel, name)
+	}
+
+	all, err := New().Extract(context.Background(), dir, rel)
+	if err != nil {
+		t.Fatalf("Extract: %v", err)
+	}
+
+	for _, f := range all {
+		if f.Kind != facts.KindSymbol || f.Name != "pkg/child.Child" {
+			continue
+		}
+		if !hasRel(f, facts.RelImplements, "pkg/father.Father") {
+			t.Errorf("Child missing Father inheritance target: %v", f.Relations)
+		}
+		if !hasRel(f, facts.RelImplements, "pkg/mother.Mother") {
+			t.Errorf("Child missing Mother inheritance target: %v", f.Relations)
+		}
+		return
+	}
+	t.Fatal("missing pkg/child.Child symbol")
+}
+
 func TestResolveCallTargets_External_DropsEdge(t *testing.T) {
 	fileModules := modSet("airflow-core/src/airflow/models/dag")
 	ff := []facts.Fact{