Upload data2/step22/depend_analysis.py with huggingface_hub

data2/step22/depend_analysis.py

@@ -0,0 +1,226 @@
+import traceback
+import os
+import json
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from tree_sitter import Language, Parser
+import tree_sitter_c
+import tree_sitter_cpp
+import tree_sitter_java
+import tree_sitter_go
+import tree_sitter_rust
+import tree_sitter_julia
+import tree_sitter_python
+
+########################################
+# 配置区
+########################################
+
+OUTPUT_DIR = "/home/weifengsun/tangou1/step2/step22/dataset_calls"
+SCORE_THRESHOLD = 0.92
+
+EXCLUDE_DIRS = {
+    ".git", "node_modules", "vendor", "third_party",
+    "build", "dist", "target", "__pycache__"
+}
+
+LANGUAGE_CONFIG = {
+    "python": {
+        "ext": [".py"],
+        "function_nodes": ["function_definition"],
+        "call_nodes": ["call"],
+        "name_field": "name",
+    },
+    "c": {
+        "ext": [".c", ".h"],
+        "function_nodes": ["function_definition"],
+        "call_nodes": ["call_expression"],
+        "name_field": "declarator",
+    },
+    "cpp": {
+        "ext": [".cpp", ".cc", ".cxx", ".hpp", ".hh"],
+        "function_nodes": ["function_definition"],
+        "call_nodes": ["call_expression"],
+        "name_field": "declarator",
+    },
+    "java": {
+        "ext": [".java"],
+        "function_nodes": ["method_declaration", "constructor_declaration"],
+        "call_nodes": ["method_invocation", "call_expression"],
+        "name_field": "name",
+    },
+    "go": {
+        "ext": [".go"],
+        "function_nodes": ["function_declaration", "method_declaration"],
+        "call_nodes": ["call_expression"],
+        "name_field": "name",
+    },
+    "rust": {
+        "ext": [".rs"],
+        "function_nodes": ["function_item"],
+        "call_nodes": ["call_expression"],
+        "name_field": "name",
+    },
+    "julia": {
+        "ext": [".jl"],
+        "function_nodes": ["function_definition"],
+        "call_nodes": ["call_expression"],
+        "name_field": "name",
+    },
+}
+
+EXT_TO_LANG = {}
+for lang, cfg in LANGUAGE_CONFIG.items():
+    for e in cfg["ext"]:
+        EXT_TO_LANG[e] = lang
+
+########################################
+# worker 初始化
+########################################
+
+LANGUAGES = {
+    "python": Language(tree_sitter_python.language()),
+    "c": Language(tree_sitter_c.language()),
+    "cpp": Language(tree_sitter_cpp.language()),
+    "java": Language(tree_sitter_java.language()),
+    "go": Language(tree_sitter_go.language()),
+    "rust": Language(tree_sitter_rust.language()),
+    "julia": Language(tree_sitter_julia.language()),
+}
+
+def init_worker():
+    global PARSERS
+    PARSERS = {}
+    for lang in LANGUAGE_CONFIG:
+        try:
+            parser = Parser()
+            parser.set_language(LANGUAGES[lang])
+            PARSERS[lang] = parser
+        except Exception:
+            print(f"Failed to load parser for {lang}")
+            print(traceback.format_exc())
+            pass
+
+########################################
+# jsonl 读取
+########################################
+
+def load_high_score_functions(jsonl_path, score_threshold=SCORE_THRESHOLD):
+    funcs = []
+    names = set()
+    with open(jsonl_path, "r", encoding="utf-8") as f:
+        for line in f:
+            obj = json.loads(line)
+            if obj.get("score", 0) >= score_threshold:
+                funcs.append(obj)
+                names.add(obj["name"])
+    return funcs, names
+
+########################################
+# Tree-sitter 提取调用
+########################################
+
+def extract_calls_from_node(node, src, call_types):
+    calls = set()
+    def walk(n):
+        if n.type in call_types:
+            for c in n.children:
+                if c.type == "identifier":
+                    calls.add(src[c.start_byte:c.end_byte].decode())
+        for c in n.children:
+            walk(c)
+    walk(node)
+    return calls
+
+def find_function_node_by_line(tree, start_line, func_nodes):
+    result = None
+    def walk(node):
+        nonlocal result
+        if result is not None:
+            return
+        if node.type in func_nodes and node.start_point[0] + 1 == start_line:
+            result = node
+            return
+        for c in node.children:
+            walk(c)
+    walk(tree.root_node)
+    return result
+
+########################################
+# 匹配函数名
+########################################
+
+def match_names(called_names, indexed_names):
+    matches = {}
+    for call in called_names:
+        hits = [name for name in indexed_names if call == name or name.startswith(call) or call.startswith(name)]
+        if hits:
+            matches[call] = hits
+    return matches
+
+########################################
+# 单项目分析
+########################################
+
+def process_project_calls(project_path):
+    project_name = os.path.basename(project_path.rstrip("/"))
+    input_jsonl = os.path.join(project_path, "functions.jsonl")  # 假设每项目jsonl在项目目录下
+    output_path = os.path.join(OUTPUT_DIR, project_name, "calls.jsonl")
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    funcs, name_set = load_high_score_functions(input_jsonl)
+
+    with open(output_path, "w", encoding="utf-8") as out:
+        for f in funcs:
+            file_path = f["file"]
+            lang = f["language"]
+            if not lang or lang not in PARSERS:
+                continue
+            try:
+                src = open(file_path, "rb").read()
+                parser = PARSERS[lang]
+                tree = parser.parse(src)
+                node = find_function_node_by_line(tree, f["start_line"], LANGUAGE_CONFIG[lang]["function_nodes"])
+                if not node:
+                    continue
+                calls = extract_calls_from_node(node, src, LANGUAGE_CONFIG[lang]["call_nodes"])
+                matched = match_names(calls, name_set)
+                out.write(json.dumps({
+                    "function": f["name"],
+                    "qualified_name": f.get("qualified_name", f["name"]),
+                    "calls": sorted(list(calls)),
+                    "matched": matched,
+                    "file": f["file"],
+                    "start_line": f["start_line"],
+                    "end_line": f["end_line"],
+                    "score": f.get("score", 0),
+                    "language": lang
+                }, ensure_ascii=False) + "\n")
+            except Exception:
+                continue
+    return project_name
+
+########################################
+# 主入口
+########################################
+
+def load_projects(root):
+    return [os.path.join(root, d) for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]
+
+def main():
+    projects_root = "/home/weifengsun/tangou1/domain_code/src/workdir/repos_filtered"
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+    projects = load_projects(projects_root)
+
+    with ProcessPoolExecutor(max_workers=min(os.cpu_count(), 32), initializer=init_worker) as pool:
+        futures = {pool.submit(process_project_calls, p): p for p in projects}
+        for f in as_completed(futures):
+            proj = futures[f]
+            try:
+                f.result()
+                print(f"[OK] {proj}")
+            except Exception as e:
+                print(f"[FAIL] {proj}: {e}")
+
+if __name__ == "__main__":
+    main()