rust语言功能增强

test/service/handler.py

@@ -5,11 +5,12 @@ import traceback
 
 TARGET_LANGUAGE = set(["python", "cpp", "go", "rust", "node", "javascript"])
 PROJECT_PATH = "/home/rocker/project"
-BUILD_TIMEOUT_S = 5
+BUILD_TIMEOUT_S = 3
+BASE_CODE_PATH = "/tmp/code"
 
 def go_handler():
     source_path = f"{PROJECT_PATH}/go_project/main.go"
-    subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
+    subprocess.run(["cp", BASE_CODE_PATH, source_path], capture_output=True)
     main_path = f"{PROJECT_PATH}/go_project/main"
     ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "go", "build", "-o", main_path], cwd=f"{PROJECT_PATH}/go_project", capture_output=True)
     std_err = ret.stderr.decode("utf-8")
@@ -19,7 +20,7 @@ def go_handler():
 
 def cpp_handler():
     source_path = f"{PROJECT_PATH}/cpp_project/main.cpp"
-    subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
+    subprocess.run(["cp", BASE_CODE_PATH, source_path], capture_output=True)
     main_path = f"{PROJECT_PATH}/cpp_project/main"
     ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "g++", source_path, "-o", main_path], capture_output=True)
     std_err = ret.stderr.decode("utf-8")
@@ -29,12 +30,25 @@ def cpp_handler():
 
 def node_handler():
     main_path = f"{PROJECT_PATH}/node_project/main.js"
-    subprocess.run(["cp", "/tmp/code", main_path], capture_output=True)
+    subprocess.run(["cp", BASE_CODE_PATH, main_path], capture_output=True)
     return "", main_path
 
 def rust_handler():
+    lines = []
+    is_have_main = False
+    with open(BASE_CODE_PATH, "r") as f:
+        for line in f.readlines():
+            lines.append(line)
+            if "fn" in line and "main" in line:
+                is_have_main = True
+    if is_have_main == False:
+        lines.insert(0, "fn main() {\n")
+        lines.append("}\n")
+    
     source_path = f"{PROJECT_PATH}/rust_project/src/main.rs"
-    subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
+    with open(source_path, "w") as f:
+        f.writelines(lines)
+
     ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "cargo", "build"], cwd=f"{PROJECT_PATH}/rust_project", capture_output=True)
     std_err = ret.stderr.decode("utf-8")
     if len(std_err) > 0 and ret.returncode != 0:
@@ -45,7 +59,7 @@ def rust_handler():
 
 def python_handler():
     main_path = f"{PROJECT_PATH}/python_project/main.py"
-    subprocess.run(["cp", "/tmp/code", main_path], capture_output=True)
+    subprocess.run(["cp", BASE_CODE_PATH, main_path], capture_output=True)
     return "", main_path
 
 
@@ -98,3 +112,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+