apply transformations on the original unmodified fuzz targets.

Author: joyguoguo

fuzz/collect_fuzz_python.py

@@ -198,7 +198,15 @@ def transform_repos(repos: list[str], jobs: int):
     def _transform_repo(repo: str):
         project_name = os.path.basename(repo)
         oss_fuzz_dir = Path(repo).parent.parent
-        targets = discover_targets(project_name, oss_fuzz_dir)
+        raw_targets = discover_targets(project_name, oss_fuzz_dir)
+        
+        # 只需移除目标名称中的 "_print1"，不要添加任何新后缀
+        transformed_targets = [t.replace("_print1", "") for t in raw_targets]
+        
+        # 去重
+        targets = list(set(transformed_targets))
+        
+        # 传递给 generate_test_template 的是简单目标名称
         return [generate_test_template(t, repo) for t in targets]
 
     with ProcessingPool(jobs) as p:
@@ -209,7 +217,10 @@ def generate_test_template(target_name: str, repo_path: str):
     """
     Generate Python test template using AST for more precise code transformations
     """
-    src_file = pjoin(repo_path, target_name + ".py")
+    src_file = pjoin(repo_path, target_name)
+    logging.info(f"Generating test template for {src_file}")    
+    if not src_file.endswith(".py"):
+        src_file += ".py"    
     if not os.path.exists(src_file):
         logging.error(f"Source target file not found: {src_file}")
         return None
@@ -253,7 +264,7 @@ def generate_test_template(target_name: str, repo_path: str):
         with open(init_path, "w", encoding="utf-8") as f:
             f.write("")
 
-    template_path = pjoin(template_dir, f"{target_name}.py")
+    template_path = pjoin(template_dir, f"{os.path.splitext(target_name)[0]}.py")
     with open(template_path, "w", encoding="utf-8") as f:
         f.write(shebang + cleaned_code.strip() + "\n")
 
@@ -294,9 +305,6 @@ def visit_FunctionDef(self, node):
             if param_name:
                 self.add_param_assignment(node, param_name)
 
-            # f. 删除所有 print(原参数名) 的语句
-            if param_name:
-                self.remove_print_param(node, param_name)
 
         # 确保继续遍历子节点
         self.generic_visit(node)
@@ -400,7 +408,7 @@ def visit_FunctionDef(self, node):
 
 def substitute_one_repo(
     repo: str,
-    targets: list[str],
+    targets: list[tuple],  # 每个元素是 (transformed_target, raw_target)
     n_fuzz: int,
     strategy: str,
     max_len: int,
@@ -414,23 +422,25 @@ def substitute_one_repo(
     template_dir = pjoin(repo, "tests-gen")
     os.makedirs(template_dir, exist_ok=True)
 
-    for target_name in targets:
-        source_file = pjoin(template_dir, f"{target_name}.py")
+    for transformed_target, raw_target in targets:
+        # 使用转换后的目标名称构建模板文件路径
+        source_file = pjoin(template_dir, transformed_target + ".py")
+        
+        # 使用原始目标名称构建输入文件路径
+        input_path = pjoin(input_dir, raw_target)
+        
+        # 确保源文件存在
         if not os.path.exists(source_file):
             logging.warning(f"Source file not found: {source_file}")
             continue
-
-        input_path = pjoin(input_dir, f"{target_name}")
         if not os.path.exists(input_path):
             logging.warning(f"Input file not found: {input_path}")
             continue
-
+        
         # 读取所有有效的输入数据
         valid_inputs = []
-        # 首先读取文件内容，然后关闭文件
         with open(input_path, "rb") as f_input:
             lines = f_input.readlines()
-
             # 文件已关闭，现在处理数据
             for line in lines:
                 # 使用 errors='replace' 确保解码不会失败
@@ -452,10 +462,12 @@ def substitute_one_repo(
                     valid_inputs.append(line)
 
         if not valid_inputs:
-            logging.warning(f"No valid inputs found for {target_name}")
+            # 使用 transformed_target 而不是 target_name
+            logging.warning(f"No valid inputs found for {transformed_target}")
             continue
 
-        logging.info(f"Loaded {len(valid_inputs)} inputs for {target_name}")
+        # 使用 transformed_target 而不是 target_name
+        logging.info(f"Loaded {len(valid_inputs)} inputs for {transformed_target}")
         # 策略选择输入
         if strategy == "shuffle":
             random.shuffle(valid_inputs)
@@ -487,28 +499,22 @@ def substitute_one_repo(
                 # 生成新代码
                 new_code = astunparse.unparse(new_tree)
 
-                out_path = pjoin(template_dir, f"{target_name}.testgen_{idx}.py")
+                # 使用 transformed_target 而不是 target_name
+                out_path = pjoin(template_dir, f"{transformed_target}.testgen_{idx}.py")
                 with open(out_path, "w") as f_out:
                     f_out.write(new_code)
 
-            # 格式化代码
-            formatter_installed = True
-            try:
-                subprocess.run(["black", out_path], 
-                            check=False, 
-                            stdout=subprocess.DEVNULL,  # 隐藏输出
-                            stderr=subprocess.DEVNULL)  # 隐藏错误
-            except FileNotFoundError:
-                if formatter_installed:  # 避免多次记录
-                    logging.warning("Black code formatter not found. For better formatting, install with:")
-                    logging.warning("pip install black")
-                    formatter_installed = False
+                # 格式化代码
+                try:
+                    subprocess.run(["black", out_path], check=False)
+                except FileNotFoundError:
+                    logging.warning("Black formatter not found, skipping formatting")
 
             except SyntaxError as e:
                 logging.error(f"Syntax error when processing {source_file}: {e}")
             except Exception as e:
-                logging.error(f"Error generating test case for {target_name}: {e}")
-
+                # 使用 transformed_target 而不是 target_name
+                logging.error(f"Error generating test case for {transformed_target}: {e}")
 
 def testgen_repos(
     repos: list[str],
@@ -529,28 +535,34 @@ def testgen_repos(
         max_len (int): Maximum length
         sim_thresh (float): Similarity threshold
     """
-    # First get all targets
-    targets_list = []
+    # First get all targets and apply transformation
+    target_map = {}
     for repo in repos:
         project_name = os.path.basename(repo)
         oss_fuzz_dir = Path(repo).parent.parent
-        targets = discover_targets(project_name, oss_fuzz_dir)
-        targets_list.append(targets)
-
-    target_map = {repo: targets for repo, targets in zip(repos, targets_list)}
+        raw_targets = discover_targets(project_name, oss_fuzz_dir)
+        
+        # 保存原始目标名称和转换后的目标名称
+        transformed_targets = [t.replace("_print1", "") for t in raw_targets]
+        targets = list(zip(transformed_targets, raw_targets))  # (转换后, 原始)
+        target_map[repo] = targets
 
     # Process each repository in parallel
     with ProcessingPool(jobs) as p:
         list(
             p.map(
                 lambda item: substitute_one_repo(
-                    item[0], item[1], n_fuzz, strategy, max_len, sim_thresh
+                    item[0],         # repo path
+                    item[1],         # list of (transformed, raw) targets
+                    n_fuzz, 
+                    strategy, 
+                    max_len, 
+                    sim_thresh
                 ),
                 target_map.items(),
             )
         )
 
-
 def main(
     repo_id: str = "data/valid_projects.txt",
     repo_root: str = "fuzz/oss-fuzz/projects/",