Update app/main.py

app/main.py

@@ -1,16 +1,15 @@
 import io
 import os
 import traceback
-
+import zipfile
 from fastapi import FastAPI, HTTPException, status
 from fastapi.concurrency import run_in_threadpool
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from opentrons.simulate import simulate, format_runlog
 
-
 app = FastAPI()
 
-
 class Protocol(BaseModel):
     name: str
     content: str
@@ -19,114 +18,163 @@ class Protocol(BaseModel):
 async def root():
     return {"message": "Opentrons simulation API is running."}
 
-# no file store in memory, just simulate the protocol
+@app.get("/protocols")
+async def list_saved_protocols():
+    """
+    storageディレクトリに保存されているプロトコルファイルの一覧を返します。
+    """
+    storage_dir = "storage"
+    if not os.path.isdir(storage_dir):
+        # ディレクトリが存在しない場合は空のリストを返す
+        return {"protocols": []}
+    
+    try:
+        # os.listdirはブロッキングI/Oなのでスレッドプールで実行
+        file_list = await run_in_threadpool(os.listdir, storage_dir)
+        return {"protocols": sorted(file_list)}
+    except Exception as e:
+        print(f"Error listing files: {e}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="An error occurred while listing the protocols."
+        )
+
+@app.get("/protocols/download")
+async def download_all_protocols():
+    """
+    storageディレクトリ内のすべてのプロトコルをzipファイルにまとめてダウンロードします。
+    """
+    storage_dir = "storage"
+    
+    if not os.path.isdir(storage_dir) or not os.listdir(storage_dir):
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="No protocols found to download."
+        )
+
+    # メモリ上でzipファイルを作成
+    zip_io = io.BytesIO()
+    
+    with zipfile.ZipFile(zip_io, mode='w', compression=zipfile.ZIP_DEFLATED) as temp_zip:
+        # run_in_threadpoolの中でファイルI/Oを実行
+        def create_zip_in_thread():
+            for root, _, files in os.walk(storage_dir):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    # zipファイル内でのパスをルートからの相対パスに設定
+                    temp_zip.write(file_path, os.path.relpath(file_path, storage_dir))
+        
+        await run_in_threadpool(create_zip_in_thread)
+
+    # BytesIOのポインタを先頭に戻す
+    zip_io.seek(0)
+
+    return StreamingResponse(
+        content=zip_io,
+        media_type="application/zip",
+        headers={"Content-Disposition": "attachment; filename=protocols.zip"}
+    )
+
 @app.post("/simulate")
 async def simulate_protocol(protocol: Protocol):
+    """
+    プロトコルの内容を受け取り、シミュレーションのみを実行します。
+    ファイルへの保存は行いません。
+    """
     try:
         protocol_file = io.StringIO(protocol.content)
+        
         run_log, _ = await run_in_threadpool(
             simulate,
             protocol_file=protocol_file,
             file_name=protocol.name
         )
-        # return execution log
-        return  {
-            'protocol_name': protocol.name,
-            'run_status': 'success',
-            'run_log': format_runlog(run_log)
+        
+        return {
+            "protocol_name": protocol.name,
+            "run_status": "success",
+            "run_log": format_runlog(run_log)
         }
     except Exception as e:
-        print(f"Error simulating protocol: {e}")
-        
+        print(f"Simulation Error: {e}")
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail={
-                "message": "Error simulating protocol",
+                "message": "Failed to simulate protocol.",
                 "error_type": type(e).__name__,
                 "error_details": str(e),
                 "traceback": traceback.format_exc()
             }
         )
 
-# @app.get("/protocols")
-# def read_protocols():
-#     protocols = get_file_names()
-#     if len(protocols) == 0:
-#         return "no stored protocol"
-#     sorted_protocols = sorted(protocols)
-#     return {"protocols": sorted_protocols}
-
-# @app.get("/protocols/{protocol_id}")
-# def read_protocol(protocol_id: int, q: str = None):
-#     # ToDo search folder and display a protocol file if there is a target protocol
-#     if q:
-#         return {"protocol_id": protocol_id, "q": q}
-#     return {"protocol_id": protocol_id}
-
-
-def get_file_names():
-    folder_path = 'storage'
-    # Check if the storage folder exists
-    if not os.path.exists(folder_path):
-        # Create the storage folder
-        os.makedirs(folder_path)
-        return []
-
-    file_names = []
-    for root, dirs, files in os.walk(folder_path):
-        for file in files:
-            file_names.append(file)
-
-    return file_names
-
-
-@app.post("/protocol")
-def upload_protocol(protocol: Protocol):
-    file_path = 'storage/' + protocol.name
-    save_result = save_text_as_file(protocol.content, file_path)
-
-    if type(save_result) == str:
-        # response = run_protocol_on_simulator(protocol.content)
-        response = call_opentrons_simulate(file_path)
-        print('response', response)
-        if response["status"] == "success":
-            return {"protocol_name": protocol.name, "run_status": "success", "run_log": response['run_log']}
-        else:
-            return {"protocol_name": protocol.name, "run_status": "failure", "error_message": response['error_message']}
-    else:
-        return {"error_message": "something wrong while saving a protocol"}
-
-
-def call_opentrons_simulate(protocol_path: str):
-    command = f"opentrons_simulate {protocol_path}"
-
-    result = subprocess.run(command, shell=True,
-                            capture_output=True, text=True)
-
-    if result.returncode == 0:
-        # print("Command executed successfully! Output:")
-        # print(result.stdout)
-        return {"status": "success", "run_log": result.stdout}
-    else:
-        print("Command failed. Error message:")
-        # print(result.stderr)
-        return {"status": "error", "error_message": result.stderr}
+@app.post("/protocols")
+async def save_and_simulate_protocol(protocol: Protocol):
+    """
+    プロトコルをファイルに保存し、続けてシミュレーションを実行します。
+    保存とシミュレーション両方の結果を返します。
+    """
+    # --- 1. プロトコルファイルを保存 ---
+    storage_dir = "storage"
+    os.makedirs(storage_dir, exist_ok=True)
+    file_path = os.path.join(storage_dir, protocol.name)
+    
+    try:
+        # スレッドプールでファイルを保存し、実際に保存されたパスを受け取る
+        saved_path = await run_in_threadpool(save_protocol_file, protocol.content, file_path)
+        save_message = f"Protocol '{os.path.basename(saved_path)}' saved successfully."
+    except Exception as e:
+        print(f"File saving error: {e}")
+        # ファイル保存に失敗した場合は、ここで処理を中断してエラーを返す
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="An error occurred while saving the file."
+        )
 
+    # --- 2. シミュレーションを実行 ---
+    try:
+        protocol_file = io.StringIO(protocol.content)
+        
+        # スレッドプールでシミュレーションを実行
+        run_log, _ = await run_in_threadpool(
+            simulate,
+            protocol_file=protocol_file,
+            file_name=protocol.name
+        )
+        
+        # 保存とシミュレーション両方の成功を返す
+        return {
+            "save_status": "success",
+            "save_message": save_message,
+            "protocol_name": protocol.name,
+            "run_status": "success",
+            "run_log": format_runlog(run_log)
+        }
+    except Exception as e:
+        # ファイル保存は成功したが、シミュレーションでエラーが発生した場合
+        print(f"Simulation Error after save: {e}")
+        # 保存が成功したことを伝えつつ、シミュレーションのエラー情報を返す
+        return {
+            "save_status": "success",
+            "save_message": save_message,
+            "protocol_name": protocol.name,
+            "run_status": "failure",
+            "error_details": str(e),
+            "traceback": traceback.format_exc()
+        }
 
-def save_text_as_file(text, file_path):
-    base = os.path.splitext(file_path)[0]
-    ext = os.path.splitext(file_path)[1]
+def save_protocol_file(content: str, file_path: str) -> str:
+    """
+    指定されたパスにテキストコンテンツを保存します。
+    ファイル名が重複する場合は、連番を付与して新しいパスに保存します。
+    実際に保存したファイルのパスを返します。
+    """
+    base, ext = os.path.splitext(file_path)
     counter = 1
-
-    while os.path.exists(file_path):
-        file_path = base + "_" + str(counter) + ext
+    new_path = file_path
+    while os.path.exists(new_path):
+        new_path = f"{base}_{counter}{ext}"
         counter += 1
-
-    try:
-        with open(file_path, 'w') as file:
-            file.write(text)
-        print("Text saved successfully as", file_path)
-        return file_path
-    except Exception as e:
-        print("An error occurred while saving the file:", str(e))
-        return False
+        
+    with open(new_path, 'w') as f:
+        f.write(content)
+    return new_path