feat: complete_code

Co1lin · Nov 13, 2024 · 83bf60e · 83bf60e
1 parent a189eb9
commit 83bf60e
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/cweval/evaluate.py b/cweval/evaluate.py
@@ -36,6 +36,7 @@
     COMPILE_DIR,
     LANGS_COMPILE,
     compile_list,
+    complete_code,
     get_code_from,
     pass_at_k,
     run_in_subprocess,
@@ -103,10 +104,7 @@ def _parse_raw_file(self, raw_file_path: str) -> str:
         tot_code = f'{raw_code}\n\n{entrypoint_src_line}\n{entrypoint_code}\n'
 
         lang = os.path.splitext(raw_file_path)[1][1:]
-        if lang == 'c':
-            tot_code = f'''#include <stdio.h>
-#include <stdbool.h>
-{tot_code}'''
+        tot_code = complete_code(tot_code, lang)
 
         return tot_code