Fix cross-file inheritance and ghost type references in AST extractor

graphify/extract.py

@@ -7820,6 +7820,71 @@ def _collect_python_symbol_resolution_facts(
                     )
                 )
 
+        # NEW: Collect heritage and annotation uses for classes/functions so
+        # they can be resolved via import/export facts in the post-pass.
+        # Walk top-level class definitions and collect base names and any
+        # annotation type names as _SymbolUseFact entries.
+        stem = _file_stem(path)
+        for node in _walk_python_tree(root_node):
+            if node.type == "class_definition":
+                # class name -> nid (file-scoped)
+
+                name_node = node.child_by_field_name("name")
+                if name_node is None:
+                    continue
+                class_name = source[name_node.start_byte:name_node.end_byte].decode("utf-8", errors="replace")
+                class_nid = _make_id(stem, class_name)
+                # Collect superclass references so they can be resolved
+                # against import/export facts during the symbol-resolution post-pass.
+
+                supers = node.child_by_field_name("superclasses")
+                if supers is not None:
+                    for s in supers.children:
+                        if s.type == "identifier":
+                            base_name = source[s.start_byte:s.end_byte].decode("utf-8", errors="replace")
+                            facts.uses.append(_SymbolUseFact(path, class_nid, base_name, "inherits", "type", s.start_point[0] + 1))
+                        elif s.type == "attribute":
+                            #qualified name like mod.Base -> tail
+                            raw = source[s.start_byte:s.end_byte].decode("utf-8", errors="replace") 
+                            tail = raw.rsplit(".", 1)[-1]
+                            facts.uses.append(_SymbolUseFact(path, class_nid, tail, "inherits", "type", s.start_point[0] + 1))
+
+            # Collect annotation type references for both classes and functions, which are handled the same way for import resolution purposes. Walk function definitions and class-level variable declarations; for any annotation node, collect identifier and attribute references as _SymbolUseFact entries with relation "references" and context "annotation".
+            if node.type == "function_definition":
+                fn_name_node = node.child_by_field_name("name")
+                if fn_name_node is None:
+                    continue
+                fn_name = source[fn_name_node.start_byte:fn_name_node.end_byte].decode("utf-8", errors = "replace")
+                fn_nid = _make_id(stem, fn_name)
+                # Parameters
+                params = node.child_by_field_name("parameters")
+                if params is not None:
+                    for p in params.children:
+                        if p.type in ("typed_parameter", "default_parameter"):
+                            type_node = p.child_by_field_name("type")
+                            if type_node is not None:
+                                # Collect simple identifiers and qualified attribute references (e.g. `mod.Type`) appearing in annotations as separate facts; the post-pass will attempt to resolve both against imports/exports.
+                                for c in type_node.children:
+                                    if c.type == "identifier":
+                                        tname = source[c.start_byte:c.end_byte].decode("utf-8", errors="replace")
+                                        facts.uses.append(_SymbolUseFact(path, fn_nid, tname, "references", "annotation", c.start_point[0] + 1))
+                                    elif c.type == "attribute":
+                                        raw = source[c.start_byte:c.end_byte].decode("utf-8", errors="replace")
+                                        tail = raw.rsplit(".", 1)[-1]
+                                        facts.uses.append(_SymbolUseFact(path, fn_nid, tail, "references", "annotation", c.start_point[0] + 1))
+
+            #return type
+            ret = node.child_by_field_name("return_type")
+            if ret is not None:
+                for c in ret.children:
+                    if c.type == "identifier":
+                        tname = source[c.start_byte:c.end_byte].decode("utf-8", errors="replace") 
+                        facts.uses.append(_SymbolUseFact(path, fn_nid, tname, "references", "annotation", c.start_point[0] + 1))
+                    elif c.type == "attribute":
+                        raw = source[c.start_byte:c.end_byte].decode("utf-8", errors="replace")
+                        tail = raw.rsplit(".", 1)[-1]
+                        facts.uses.append(_SymbolUseFact(path, fn_nid, tail, "references", "annotation", c.start_point[0] + 1))
+
 
 def _augment_symbol_resolution_edges(
     paths: list[Path],