fix data collection code

Author: Lut-hub

testexplora/build_benchmark/build_dependency_graph.py

@@ -9,6 +9,7 @@
 import networkx as nx
 from matplotlib.lines import Line2D
 import networkx as nx
+import logging
 
 VERSION = 'v2.3'
 NODE_TYPE_DIRECTORY = 'directory'
@@ -33,14 +34,18 @@
     EDGE_TYPE_IMPORTS,
 ]
 
-SKIP_DIRS = ['.github', '.git']
+SKIP_DIRS = [
+    '.git', '.github', '__pycache__', '.mypy_cache', '.pytest_cache', '.tox',
+    'build', 'dist', '.eggs', 'node_modules', '.venv', 'venv', 'env',
+    '.idea', '.vscode', '.direnv', '.pytest', '.coverage', 'docs'
+]
+
 
 
 def is_skip_dir(dirname):
-    for skip_dir in SKIP_DIRS:
-        if skip_dir in dirname:
-            return True
-    return False
+    parts = set(dirname.split(os.sep))
+    return any(skip in parts for skip in SKIP_DIRS)
+
 
 
 def handle_edge_cases(code):
@@ -124,8 +129,9 @@ def find_imports(filepath, repo_path, tree=None):
 
 
 class CodeAnalyzer(ast.NodeVisitor):
-    def __init__(self, filename):
+    def __init__(self, filename, code):
         self.filename = filename
+        self.code = code
         self.nodes = []
         self.node_name_stack = []
         self.node_type_stack = []
@@ -184,12 +190,12 @@ def _visit_func(self, node):
         self.node_type_stack.pop()
 
     def _get_source_segment(self, node):
-        with open(self.filename, 'r') as file:
-            source_code = file.read()
-        return ast.get_source_segment(source_code, node)
+        # with open(self.filename, 'r') as file:
+        #     source_code = file.read()
+        return ast.get_source_segment(self.code, node)
 
 
-# Parse the specified file, using CodeAnalyzer to analyze classes and top-level functions
+# 解析指定文件，使用CodeAnalyzer分析文件中的类和顶级函数
 def analyze_file(filepath):
     with open(filepath, 'r') as file:
         code = file.read()
@@ -198,7 +204,7 @@ def analyze_file(filepath):
             tree = ast.parse(code, filename=filepath)
         except:
             raise SyntaxError
-    analyzer = CodeAnalyzer(filepath)
+    analyzer = CodeAnalyzer(filepath, code)
     try:
         analyzer.visit(tree)
     except RecursionError:
@@ -331,7 +337,7 @@ def resolve_symlink(file_path):
 #     return file_content
 
 
-# Traverse all Python files under repo_path, building a dependency graph of files, classes, and functions
+# 遍历repo_path下的所有Python文件，构建文件、类和函数的依赖关系图
 def build_graph(repo_path, fuzzy_search=True, global_import=False):
     graph = nx.MultiDiGraph()
     file_nodes = {}
@@ -386,7 +392,7 @@ def build_graph(repo_path, fuzzy_search=True, global_import=False):
                     else:
                         with open(file_path, 'r') as f:
                             file_content = f.read()
-
+                    logging.debug(f'Analyzing file: {file_path}')
                     graph.add_node(filename, type=NODE_TYPE_FILE, code=file_content)
                     file_nodes[filename] = file_path
 
@@ -516,10 +522,9 @@ def build_graph(repo_path, fuzzy_search=True, global_import=False):
 
 def build_json_graph(repo_path, fuzzy_search=True, global_import=False):
     graph = build_graph(repo_path, fuzzy_search, global_import)
-    temp_graph_dict = nx.node_link_data(graph, edges="links")
+    temp_graph_dict = nx.node_link_data(graph)
     codes = {}
     folders = {}
-
     nodes = temp_graph_dict['nodes']
     edges = temp_graph_dict['links']
     edge_dic = {}
@@ -718,9 +723,9 @@ def process_inheritance_node(_inheritance_node):
 
                     for sub_node in ast.walk(body_item):
                         if isinstance(sub_node, ast.Call):
-                            if isinstance(sub_node.func, ast.Name):  # plain function or class
+                            if isinstance(sub_node.func, ast.Name):  # 普通函数或类
                                 add_invoke(sub_node.func.id)
-                            if isinstance(sub_node.func, ast.Attribute):  # member function
+                            if isinstance(sub_node.func, ast.Attribute):  # 成员函数
                                 add_invoke(sub_node.func.attr)
                     break
             break
@@ -732,7 +737,7 @@ def analyze_invokes(node, code_tree, graph, repo_path):
     caller_name = node.split(':')[-1].split('.')[-1]
     file_path = os.path.join(repo_path, node.split(':')[0])
 
-    # Store found call relationships
+    # 存储找到的调用关系
     invocations = []
 
     def add_invoke(func_name):
@@ -764,7 +769,7 @@ def traverse_call(_node):
             # Recursively traverse child nodes
             traverse_call(child)
 
-    # Traverse AST nodes to find call relationships
+    # 遍历 AST 节点以找到调用关系
     for ast_node in ast.walk(code_tree):
         if (
             isinstance(ast_node, (ast.FunctionDef, ast.AsyncFunctionDef))
@@ -774,11 +779,11 @@ def traverse_call(_node):
             imports = find_imports(file_path, repo_path, tree=ast_node)
             add_imports(node, imports, graph, repo_path)
 
-            # Traverse function decorators
+            # 遍历函数装饰器
             for decorator_node in ast_node.decorator_list:
                 process_decorator_node(decorator_node)
 
-            # Traverse all invoke child nodes in the function body (excluding inner functions and classes)
+            # 遍历函数体内的所有invoke子节点 (不包括内部函数、类)
             traverse_call(ast_node)
             break
 
@@ -949,12 +954,10 @@ def main():
 
 
 if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        '--repo_path', type=str, default='DATA/repo/pallets__flask-5063'
+    logging.basicConfig(level=logging.DEBUG)
+    a = build_json_graph(
+        '/home/superbench/jiaxiang/Dev-TestCaseGen/build_benchmark/temp_repos/pyramid',
+        fuzzy_search=True
     )
-    parser.add_argument('--visualize', action='store_true')
-    parser.add_argument('--global_import', action='store_true')
-    args = parser.parse_args()
 
-    main()
+    b = 1

testexplora/build_benchmark/parse_repo.py

@@ -1,6 +1,8 @@
 import ast
 import asttokens
 import os
+import json
+import re
 
 def extract_classes_and_functions(filepath, repo_saved_path):
     file_text = ""
@@ -41,6 +43,8 @@ class ClassVisitor(ast.NodeVisitor):
         def visit_ClassDef(self, node):
             try:
                 class_docstring = ast.get_docstring(node)
+                if class_docstring is None:
+                    class_docstring = ""
                 class_docstring_end = get_docstring_end_line(node)
                 clazz = {
                     "name": node.name,
@@ -51,9 +55,12 @@ def visit_ClassDef(self, node):
                     "docstring": class_docstring if class_docstring is not None else "",
                     "methods": []
                 }
-                all_lines[prefix + ":" + node.name] = {"line": (clazz["start_line"], clazz["end_line"]), "docstring": class_docstring, "docstring_end_line": clazz["docstring_end_line"], "file": prefix}
+                all_lines[prefix + ":" + node.name] = {"line": (clazz["start_line"], clazz["end_line"]), "docstring": class_docstring, "docstring_end_line": clazz["docstring_end_line"], "file": prefix, "class": "", "type": "class", "methods": [], "parameters": {}, "returns": {}}
+                all_methods = []
+                all_parameters = {}
+                all_returns = {}    
                 for item in node.body:
-                    if isinstance(item, ast.FunctionDef):
+                    if isinstance(item, ast.FunctionDef) or isinstance(item, ast.AsyncFunctionDef):
                         returns = [
                             ast.unparse(stmt.value) if stmt.value else None
                             for stmt in item.body if isinstance(stmt, ast.Return)
@@ -62,19 +69,29 @@ def visit_ClassDef(self, node):
                         docstring_end_line = get_docstring_end_line(item)
                         decorators = [ast.unparse(d) for d in item.decorator_list]
                         body_text = atok.get_text(item)[item.body[0].col_offset:] if item.body else ""
+                        parmeters = {}
+                        for arg in item.args.args:
+                            annotation = ast.unparse(arg.annotation) if arg.annotation else None
+                            parmeters[arg.arg] = annotation
 
                         method = {
                             "name": item.name,
                             "start_line": item.decorator_list[0].lineno if item.decorator_list else item.lineno,
                             "end_line": item.end_lineno,
                             "docstring_end_line": docstring_end_line,
-                            "parameters": [arg.arg for arg in item.args.args],
+                            "parameters": parmeters,
                             "returns": returns if returns else None,
                             "docstring": docstring if docstring is not None else "",
                             "decorators": decorators if decorators else None,
                             "body": body_text.strip()
                         }
-                        all_lines[prefix + ":" + node.name + "." + item.name] = {"line": (method["start_line"], method["end_line"]), "docstring": method["docstring"], "docstring_end_line": method["docstring_end_line"], "file": prefix}
+                        all_lines[prefix + ":" + node.name + "." + item.name] = {"line": (method["start_line"], method["end_line"]), "docstring": method["docstring"], "docstring_end_line": method["docstring_end_line"], "file": prefix, "class": prefix + ":" + node.name, "type": "method", "methods": [], "parameters": method["parameters"], "returns": method["returns"]}
+                        all_methods.append(prefix + ":" + node.name + "." + item.name)
+                        all_parameters[prefix + ":" + node.name + "." + item.name] = method["parameters"]
+                        all_returns[prefix + ":" + node.name + "." + item.name] = method["returns"]
+                all_lines[prefix + ":" + node.name]["methods"] = all_methods
+                all_lines[prefix + ":" + node.name]["parameters"] = all_parameters
+                all_lines[prefix + ":" + node.name]["returns"] = all_returns
 
                 #  classes.append(clazz)
             except Exception as e:
@@ -104,19 +121,22 @@ def visit_FunctionDef(self, node):
                     docstring = ast.get_docstring(node)
                     docstring_end_line = get_docstring_end_line(node)
                     body_text = atok.get_text(node)[node.body[0].col_offset:] if node.body else ""
-
+                    parmeters = {}
+                    for arg in node.args.args:
+                        annotation = ast.unparse(arg.annotation) if arg.annotation else None
+                        parmeters[arg.arg] = annotation
                     function = {
                         "name": node.name,
                         "start_line": node.decorator_list[0].lineno if node.decorator_list else node.lineno,
                         "end_line": node.end_lineno,
                         "docstring_end_line": docstring_end_line,
-                        "parameters": [arg.arg for arg in node.args.args],
+                        "parameters": parmeters,
                         "returns": returns if returns else None,
                         "docstring": docstring if docstring is not None else "",
                         "decorators": decorators if decorators else None,
                         "body": body_text.strip()
                     }
-                    all_lines[prefix + ":" + node.name] = {"line": (function["start_line"], function["end_line"]), "docstring": function["docstring"], "docstring_end_line": function["docstring_end_line"], "file": prefix}
+                    all_lines[prefix + ":" + node.name] = {"line": (function["start_line"], function["end_line"]), "docstring": function["docstring"], "docstring_end_line": function["docstring_end_line"], "file": prefix, "class": "", "type": "function", "methods": [], "parameters": {prefix + ":" + node.name: function["parameters"]}, "returns": {prefix + ":" + node.name: function["returns"]}}
                 except Exception as e:
                     print(f"Error while processing function {node.name}: {e}")
                 finally:
@@ -140,4 +160,73 @@ def extract_classes_and_functions_from_repo(repo_saved_path):
                 all_code_lines.update(file_lines)
             except Exception as e:
                 print(f"Error processing file {filepath}: {e}")
-    return all_code_lines
+    return all_code_lines
+
+def extract_import_lines(file_content: str):
+    """
+    Extracts import lines from the given file content: import, import as, from, and from ... import.
+    Use ast to parse the file content and extract import statements.
+    """
+    imports = []
+    try:
+        tree = ast.parse(file_content)
+    except SyntaxError as e:
+        import_pattern = re.compile(r'^\s*(import|from)\s+.*$', re.MULTILINE)
+        lines = file_content.split("\n")
+        for line in lines:
+            if import_pattern.match(line):
+                imports.append(line.strip())
+        return "\n".join(imports)
+
+    # 遍历 AST 的顶层节点
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                if alias.asname:
+                    imports.append(f"import {alias.name} as {alias.asname}")
+                else:
+                    imports.append(f"import {alias.name}")
+        elif isinstance(node, ast.ImportFrom):
+            module = node.module if node.module else "" # 'from . import ...' 的情况
+            names = []
+            for alias in node.names:
+                if alias.asname:
+                    names.append(f"{alias.name} as {alias.asname}")
+                else:
+                    names.append(alias.name)
+            if node.level > 0: # 处理相对导入，例如 from .module import name
+                module_prefix = "." * node.level
+                module = module_prefix + module
+            imports.append(f"from {module} import {', '.join(names)}")
+    return "\n".join(imports)
+
+class Repo_file_container:
+    def __init__(self, repo_saved_path: str):
+        self.repo_saved_path = repo_saved_path
+        self.file_dic = {}
+        self.get_all_files()
+    
+    def get_all_files(self):
+        """
+        Returns a list of all Python files in the dataset directory.
+        """
+        for root, _, files in os.walk(self.repo_saved_path):
+            for file in files:
+                if not file.endswith(".py"):
+                    continue
+                file_path = os.path.join(root, file)
+                try:
+                    with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+                        content = f.read()
+                    f.close()
+                except Exception as e:
+                    logging.info(f"Error reading {file_path}: {e}")
+                    continue
+                import_lines = extract_import_lines(content)
+                prefix = file_path.replace(self.repo_saved_path, "").lstrip("/")
+                self.file_dic[prefix] = {"lines": content.split("\n"), "imports": import_lines}
+    def get_file_content(self, file_key: str):
+        """
+        Returns the content of a file given its key.
+        """
+        return self.file_dic[file_key]