content: import argparse
  import json
  import os
  from typing import Any, Dict, List
  
  from rich.progress import track
  
  from evalplus.eval.utils import swallow_io
  from evalplus.evaluate import evaluate
  from tools.tsr.utils import (
      clean,
      execute_cmd,
      get_cmd_output,
      get_problems,
      get_task_ids,
      to_path,
  )
  
  
  def prepare_mutants(mutation_dir: str, dataset: str):
      pwd = os.getcwd()
      task_ids = get_task_ids(dataset)
      problems = get_problems(dataset)
      os.makedirs(mutation_dir, exist_ok=True)
      for task_id in track(task_ids, "Generating mutants"):
          task_dir = os.path.join(mutation_dir, to_path(task_id))
          os.makedirs(task_dir, exist_ok=True)
          if any(map(lambda filename: filename.startswith("m"), os.listdir(task_dir))):
              # already have mutants
              continue
          # Make groundtruth
          groundtruth_code = (
              problems[task_id]["prompt"] + problems[task_id]["canonical_solution"]
          )
          with open(os.path.join(task_dir, "gt.py"), "w") as f:
              f.write(groundtruth_code)
          # Make dummy pytest
          with open(os.path.join(task_dir, "test_dummy.py"), "w") as f:
              f.write("def test_dummy():\n    pass")
          # Use mutmut to generate mutants
          os.chdir(task_dir)
          clean(".mutmut-cache")
          execute_cmd(["mutmut run", "--paths-to-mutate=gt.py", "1>/dev/null"])
          try:
              # Collect metainfo
              total_mutants = int(
                  get_cmd_output(["mutmut", "results"]).split("\n")[-2].split("-")[-1]
              )
          except:
              total_mutants = 0
          # Dump mutants
          for i in range(1, total_mutants + 1):
              execute_cmd(["cp", "gt.py", "gt_copy.py"])
              execute_cmd(["mutmut", "apply", str(i)])
              execute_cmd(["mv", "gt.py", f"m{i}.py"])
              execute_cmd(["mv", "gt_copy.py", "gt.py"])
          # Remove gt and dummy pytest
          execute_cmd(["rm", "gt.py"])
          execute_cmd(["rm", "test_dummy.py"])
          clean(".mutmut-cache")
          os.chdir(pwd)
  
  
  def mutants_eval(mutation_dir: str, dataset: str):
      args = argparse.Namespace(
          dataset=dataset,
          samples=mutation_dir,
          base_only=False,
          parallel=None,
          i_just_wanna_run=False,
          test_details=True,
          min_time_limit=0.2,
          gt_time_limit_factor=4.0,
          mini=False,
      )
      print("Evaluating mutants... ", end="", flush=True)
      with swallow_io():
          evaluate(args)
      print("Done")
  
  
  def collect_mutation_info(
      eval_path: str, dataset: str
  ) -> Dict[str, Dict[str, List[Any]]]:
      mutation_info = {task_id: {} for task_id in get_task_ids(dataset)}
      assert os.path.isfile(
          eval_path
      ), f"mutation testing result file {eval_path} missing!"
      eval_res = json.load(open(eval_path, "r"))["eval"]
      for task_id, v in eval_res.items():
          for i_code, (status, res_list) in enumerate(v["plus"]):
              if status == "success":
                  continue
              for i_test, res in enumerate(res_list):
                  test_id = f"plus_{i_test}"
                  if res == False:
                      mutation_info[task_id].setdefault(test_id, []).append(
                          ("mutant", i_code)
                      )
      return mutation_info
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser()
      parser.add_argument("--dataset", type=str, choices=["humaneval", "mbpp"])
      parser.add_argument("--report_dir", required=True, type=str)
      args = parser.parse_args()
  
      mutation_dir = os.path.join(args.report_dir, "mutation_cache")
      prepare_mutants(mutation_dir, args.dataset)
      mutants_eval(mutation_dir, args.dataset)
      collect_mutation_info(os.path.join(mutation_dir, "eval_results.json"), args.dataset)
  