Update main.py

main.py

@@ -1,24 +1,25 @@
 from fastapi import FastAPI
+from pydantic import BaseModel
+
 import subprocess
-import re
+from regex import find_imports
+class Item(BaseModel):
+    code: str
+    filename: str = ''
 
-def find_imports(code):
-    pattern = r'Python\.import_module\("([^"]*)"\)'
-    matches = re.findall(pattern, code)
-    return matches
 
 
 app = FastAPI()
 
 @app.post("/code")
-def run_mojo_code(code:str, filename:str) -> dict:
+def run_mojo_code(item:Item):
     try:
-        imports = find_imports(code)
+        imports = find_imports(item.code)
         for imported in imports:
             subprocess.call(["python3", "-m", "pip", "install", imported], shell=True)
-        with open(filename, "w") as f:
-            f.write(code) 
+        with open(item.filename, "w") as f:
+            f.write(item.code) 
             
-        return {"sucess":True, "output": subprocess.check_output(["mojo", filename]).decode("utf-8")}, 200
+        return {"sucess":True, "output": subprocess.check_output(["mojo", item.filename]).decode("utf-8")}, 200
     except:
         return {"sucess":False}, 500