wehub-resource-sync
397 lines
18 KiB
Python
397 lines
18 KiB
Python
"""Go extractor. Moved verbatim from graphify/extract.py."""
|
|
from __future__ import annotations
|
|
|
|
|
|
from pathlib import Path
|
|
from graphify.extractors.base import _LANGUAGE_BUILTIN_GLOBALS, _file_stem, _make_id, _read_text
|
|
|
|
|
|
_GO_PREDECLARED_TYPES = frozenset({
|
|
"bool", "byte", "complex64", "complex128", "error", "float32", "float64",
|
|
"int", "int8", "int16", "int32", "int64", "rune", "string",
|
|
"uint", "uint8", "uint16", "uint32", "uint64", "uintptr", "any", "comparable",
|
|
})
|
|
|
|
def _go_collect_type_refs(node, source: bytes, generic: bool, out: list[tuple[str, str]]) -> None:
|
|
"""Walk a Go type expression; append (name, role) tuples."""
|
|
if node is None:
|
|
return
|
|
t = node.type
|
|
if t == "type_identifier":
|
|
text = _read_text(node, source)
|
|
if text and text not in _GO_PREDECLARED_TYPES:
|
|
out.append((text, "generic_arg" if generic else "type"))
|
|
return
|
|
if t == "qualified_type":
|
|
text = _read_text(node, source).rsplit(".", 1)[-1]
|
|
if text and text not in _GO_PREDECLARED_TYPES:
|
|
out.append((text, "generic_arg" if generic else "type"))
|
|
return
|
|
if t == "generic_type":
|
|
type_field = node.child_by_field_name("type")
|
|
if type_field is not None:
|
|
sub: list[tuple[str, str]] = []
|
|
_go_collect_type_refs(type_field, source, generic, sub)
|
|
out.extend(sub)
|
|
for c in node.children:
|
|
if c.type == "type_arguments":
|
|
for arg in c.children:
|
|
if arg.is_named:
|
|
_go_collect_type_refs(arg, source, True, out)
|
|
return
|
|
if t in ("pointer_type", "slice_type", "array_type", "map_type",
|
|
"channel_type", "parenthesized_type"):
|
|
for c in node.children:
|
|
if c.is_named:
|
|
_go_collect_type_refs(c, source, generic, out)
|
|
return
|
|
if node.is_named:
|
|
for c in node.children:
|
|
if c.is_named:
|
|
_go_collect_type_refs(c, source, generic, out)
|
|
|
|
def extract_go(path: Path) -> dict:
|
|
"""Extract functions, methods, type declarations, and imports from a .go file."""
|
|
try:
|
|
import tree_sitter_go as tsgo
|
|
from tree_sitter import Language, Parser
|
|
except ImportError:
|
|
return {"nodes": [], "edges": [], "error": "tree-sitter-go not installed"}
|
|
|
|
try:
|
|
language = Language(tsgo.language())
|
|
parser = Parser(language)
|
|
source = path.read_bytes()
|
|
tree = parser.parse(source)
|
|
root = tree.root_node
|
|
except Exception as e:
|
|
return {"nodes": [], "edges": [], "error": str(e)}
|
|
|
|
stem = _file_stem(path)
|
|
# Use directory name as package scope so methods on the same type across
|
|
# multiple files in a package share one canonical type node.
|
|
pkg_scope = path.parent.name or stem
|
|
str_path = str(path)
|
|
nodes: list[dict] = []
|
|
edges: list[dict] = []
|
|
seen_ids: set[str] = set()
|
|
function_bodies: list[tuple[str, object]] = []
|
|
go_imported_pkgs: set[str] = set() # local names of imported packages
|
|
|
|
def add_node(nid: str, label: str, line: int) -> None:
|
|
if nid not in seen_ids:
|
|
seen_ids.add(nid)
|
|
nodes.append({
|
|
"id": nid,
|
|
"label": label,
|
|
"file_type": "code",
|
|
"source_file": str_path,
|
|
"source_location": f"L{line}",
|
|
})
|
|
|
|
def add_edge(src: str, tgt: str, relation: str, line: int,
|
|
confidence: str = "EXTRACTED", weight: float = 1.0,
|
|
context: str | None = None) -> None:
|
|
edge = {
|
|
"source": src,
|
|
"target": tgt,
|
|
"relation": relation,
|
|
"confidence": confidence,
|
|
"source_file": str_path,
|
|
"source_location": f"L{line}",
|
|
"weight": weight,
|
|
}
|
|
if context:
|
|
edge["context"] = context
|
|
edges.append(edge)
|
|
|
|
file_nid = _make_id(str(path))
|
|
add_node(file_nid, path.name, 1)
|
|
|
|
def ensure_named_node(name: str, line: int) -> str:
|
|
nid = _make_id(pkg_scope, name)
|
|
if nid in seen_ids:
|
|
return nid
|
|
nid = _make_id(name)
|
|
if nid not in seen_ids:
|
|
# The name isn't declared in this file, so this is a cross-file reference
|
|
# (e.g. a type defined in another file of the package). Emit a SOURCELESS
|
|
# stub — like the inheritance-base path in the other extractors — so the
|
|
# corpus-level rewire can collapse it onto the real definition. A sourced
|
|
# stub here makes _disambiguate_colliding_node_ids bake the referencing
|
|
# file's path (with extension) into the id and blocks the rewire, which is
|
|
# the phantom-duplicate-node bug (#1402).
|
|
seen_ids.add(nid)
|
|
nodes.append({
|
|
"id": nid,
|
|
"label": name,
|
|
"file_type": "code",
|
|
"source_file": "",
|
|
"source_location": "",
|
|
"origin_file": str_path,
|
|
})
|
|
return nid
|
|
|
|
def emit_go_method_refs(func_node, func_nid: str, line: int) -> None:
|
|
params = func_node.child_by_field_name("parameters")
|
|
if params is not None:
|
|
for p in params.children:
|
|
if p.type != "parameter_declaration":
|
|
continue
|
|
type_node = p.child_by_field_name("type")
|
|
refs: list[tuple[str, str]] = []
|
|
_go_collect_type_refs(type_node, source, False, refs)
|
|
for ref_name, role in refs:
|
|
ctx = "generic_arg" if role == "generic_arg" else "parameter_type"
|
|
tgt = ensure_named_node(ref_name, line)
|
|
if tgt != func_nid:
|
|
add_edge(func_nid, tgt, "references", line, context=ctx)
|
|
result = func_node.child_by_field_name("result")
|
|
if result is not None:
|
|
if result.type == "parameter_list":
|
|
for p in result.children:
|
|
if p.type != "parameter_declaration":
|
|
continue
|
|
type_node = p.child_by_field_name("type")
|
|
if type_node is None:
|
|
for c in p.children:
|
|
if c.is_named:
|
|
type_node = c
|
|
break
|
|
refs = []
|
|
_go_collect_type_refs(type_node, source, False, refs)
|
|
for ref_name, role in refs:
|
|
ctx = "generic_arg" if role == "generic_arg" else "return_type"
|
|
tgt = ensure_named_node(ref_name, line)
|
|
if tgt != func_nid:
|
|
add_edge(func_nid, tgt, "references", line, context=ctx)
|
|
else:
|
|
refs = []
|
|
_go_collect_type_refs(result, source, False, refs)
|
|
for ref_name, role in refs:
|
|
ctx = "generic_arg" if role == "generic_arg" else "return_type"
|
|
tgt = ensure_named_node(ref_name, line)
|
|
if tgt != func_nid:
|
|
add_edge(func_nid, tgt, "references", line, context=ctx)
|
|
|
|
def walk(node) -> None:
|
|
t = node.type
|
|
|
|
if t == "function_declaration":
|
|
name_node = node.child_by_field_name("name")
|
|
if name_node:
|
|
func_name = _read_text(name_node, source)
|
|
line = node.start_point[0] + 1
|
|
func_nid = _make_id(stem, func_name)
|
|
add_node(func_nid, f"{func_name}()", line)
|
|
add_edge(file_nid, func_nid, "contains", line)
|
|
emit_go_method_refs(node, func_nid, line)
|
|
body = node.child_by_field_name("body")
|
|
if body:
|
|
function_bodies.append((func_nid, body))
|
|
return
|
|
|
|
if t == "method_declaration":
|
|
receiver = node.child_by_field_name("receiver")
|
|
receiver_type: str | None = None
|
|
if receiver:
|
|
for param in receiver.children:
|
|
if param.type == "parameter_declaration":
|
|
type_node = param.child_by_field_name("type")
|
|
if type_node:
|
|
receiver_type = _read_text(type_node, source).lstrip("*").strip()
|
|
break
|
|
name_node = node.child_by_field_name("name")
|
|
if not name_node:
|
|
return
|
|
method_name = _read_text(name_node, source)
|
|
line = node.start_point[0] + 1
|
|
|
|
if receiver_type:
|
|
parent_nid = _make_id(pkg_scope, receiver_type)
|
|
add_node(parent_nid, receiver_type, line)
|
|
method_nid = _make_id(parent_nid, method_name)
|
|
add_node(method_nid, f".{method_name}()", line)
|
|
add_edge(parent_nid, method_nid, "method", line)
|
|
else:
|
|
method_nid = _make_id(stem, method_name)
|
|
add_node(method_nid, f"{method_name}()", line)
|
|
add_edge(file_nid, method_nid, "contains", line)
|
|
|
|
emit_go_method_refs(node, method_nid, line)
|
|
body = node.child_by_field_name("body")
|
|
if body:
|
|
function_bodies.append((method_nid, body))
|
|
return
|
|
|
|
if t == "type_declaration":
|
|
for child in node.children:
|
|
if child.type != "type_spec":
|
|
continue
|
|
name_node = child.child_by_field_name("name")
|
|
if not name_node:
|
|
continue
|
|
type_name = _read_text(name_node, source)
|
|
line = child.start_point[0] + 1
|
|
type_nid = _make_id(pkg_scope, type_name)
|
|
add_node(type_nid, type_name, line)
|
|
add_edge(file_nid, type_nid, "contains", line)
|
|
# Type body: struct fields (with embeds) or interface embedding.
|
|
type_body = None
|
|
for tc in child.children:
|
|
if tc.type in ("struct_type", "interface_type"):
|
|
type_body = tc
|
|
break
|
|
if type_body is None:
|
|
continue
|
|
if type_body.type == "struct_type":
|
|
for fdl in type_body.children:
|
|
if fdl.type != "field_declaration_list":
|
|
continue
|
|
for field in fdl.children:
|
|
if field.type != "field_declaration":
|
|
continue
|
|
has_name = any(
|
|
fc.type == "field_identifier" for fc in field.children
|
|
)
|
|
type_node = field.child_by_field_name("type")
|
|
if type_node is None:
|
|
for fc in field.children:
|
|
if fc.is_named and fc.type != "field_identifier":
|
|
type_node = fc
|
|
break
|
|
refs: list[tuple[str, str]] = []
|
|
_go_collect_type_refs(type_node, source, False, refs)
|
|
for ref_name, role in refs:
|
|
tgt = ensure_named_node(ref_name, field.start_point[0] + 1)
|
|
if tgt == type_nid:
|
|
continue
|
|
if not has_name and role == "type":
|
|
add_edge(type_nid, tgt, "embeds",
|
|
field.start_point[0] + 1)
|
|
else:
|
|
ctx = "generic_arg" if role == "generic_arg" else "field"
|
|
add_edge(type_nid, tgt, "references",
|
|
field.start_point[0] + 1, context=ctx)
|
|
elif type_body.type == "interface_type":
|
|
for elem in type_body.children:
|
|
if elem.type != "type_elem":
|
|
continue
|
|
refs = []
|
|
for sub in elem.children:
|
|
if sub.is_named:
|
|
_go_collect_type_refs(sub, source, False, refs)
|
|
for ref_name, role in refs:
|
|
tgt = ensure_named_node(ref_name, elem.start_point[0] + 1)
|
|
if tgt == type_nid:
|
|
continue
|
|
if role == "type":
|
|
add_edge(type_nid, tgt, "embeds",
|
|
elem.start_point[0] + 1)
|
|
else:
|
|
add_edge(type_nid, tgt, "references",
|
|
elem.start_point[0] + 1, context="generic_arg")
|
|
return
|
|
|
|
if t == "import_declaration":
|
|
for child in node.children:
|
|
if child.type == "import_spec_list":
|
|
for spec in child.children:
|
|
if spec.type == "import_spec":
|
|
path_node = spec.child_by_field_name("path")
|
|
if path_node:
|
|
raw = _read_text(path_node, source).strip('"')
|
|
# Prefix with go_pkg_ so stdlib names (e.g. "context")
|
|
# don't collide with local files of the same basename.
|
|
tgt_nid = _make_id("go", "pkg", raw)
|
|
add_edge(file_nid, tgt_nid, "imports_from", spec.start_point[0] + 1, context="import")
|
|
# Track local name (alias or last path segment)
|
|
alias = spec.child_by_field_name("name")
|
|
local_name = _read_text(alias, source) if alias else raw.split("/")[-1]
|
|
if local_name and local_name != "_" and local_name != ".":
|
|
go_imported_pkgs.add(local_name)
|
|
elif child.type == "import_spec":
|
|
path_node = child.child_by_field_name("path")
|
|
if path_node:
|
|
raw = _read_text(path_node, source).strip('"')
|
|
tgt_nid = _make_id("go", "pkg", raw)
|
|
add_edge(file_nid, tgt_nid, "imports_from", child.start_point[0] + 1, context="import")
|
|
alias = child.child_by_field_name("name")
|
|
local_name = _read_text(alias, source) if alias else raw.split("/")[-1]
|
|
if local_name and local_name != "_" and local_name != ".":
|
|
go_imported_pkgs.add(local_name)
|
|
return
|
|
|
|
for child in node.children:
|
|
walk(child)
|
|
|
|
walk(root)
|
|
|
|
label_to_nid: dict[str, str] = {}
|
|
for n in nodes:
|
|
raw = n["label"]
|
|
normalised = raw.strip("()").lstrip(".")
|
|
label_to_nid[normalised] = n["id"]
|
|
|
|
seen_call_pairs: set[tuple[str, str]] = set()
|
|
raw_calls: list[dict] = []
|
|
|
|
def walk_calls(node, caller_nid: str) -> None:
|
|
if node.type in ("function_declaration", "method_declaration"):
|
|
return
|
|
if node.type == "call_expression":
|
|
func_node = node.child_by_field_name("function")
|
|
callee_name: str | None = None
|
|
is_member_call: bool = False
|
|
if func_node:
|
|
if func_node.type == "identifier":
|
|
callee_name = _read_text(func_node, source)
|
|
elif func_node.type == "selector_expression":
|
|
field = func_node.child_by_field_name("field")
|
|
operand = func_node.child_by_field_name("operand")
|
|
receiver_name = _read_text(operand, source) if operand else ""
|
|
# Package-qualified call (e.g. fmt.Println) → allow cross-file resolution.
|
|
# Receiver method call (e.g. s.logger.Log) → skip, no import evidence.
|
|
is_member_call = receiver_name not in go_imported_pkgs
|
|
if field:
|
|
callee_name = _read_text(field, source)
|
|
if callee_name and callee_name not in _LANGUAGE_BUILTIN_GLOBALS:
|
|
tgt_nid = label_to_nid.get(callee_name)
|
|
if tgt_nid and tgt_nid != caller_nid:
|
|
pair = (caller_nid, tgt_nid)
|
|
if pair not in seen_call_pairs:
|
|
seen_call_pairs.add(pair)
|
|
line = node.start_point[0] + 1
|
|
edges.append({
|
|
"source": caller_nid,
|
|
"target": tgt_nid,
|
|
"relation": "calls",
|
|
"context": "call",
|
|
"confidence": "EXTRACTED",
|
|
"source_file": str_path,
|
|
"source_location": f"L{line}",
|
|
"weight": 1.0,
|
|
})
|
|
elif callee_name:
|
|
raw_calls.append({
|
|
"caller_nid": caller_nid,
|
|
"callee": callee_name,
|
|
"is_member_call": is_member_call,
|
|
"source_file": str_path,
|
|
"source_location": f"L{node.start_point[0] + 1}",
|
|
})
|
|
for child in node.children:
|
|
walk_calls(child, caller_nid)
|
|
|
|
for caller_nid, body_node in function_bodies:
|
|
walk_calls(body_node, caller_nid)
|
|
|
|
valid_ids = seen_ids
|
|
clean_edges = []
|
|
for edge in edges:
|
|
src, tgt = edge["source"], edge["target"]
|
|
if src in valid_ids and (tgt in valid_ids or edge["relation"] in ("imports", "imports_from")):
|
|
clean_edges.append(edge)
|
|
|
|
return {"nodes": nodes, "edges": clean_edges, "raw_calls": raw_calls}
|