Add ExptDataWriter

- porting from arayabrain#366

studio/app/common/core/experiment/experiment_reader.py

@@ -54,25 +54,3 @@ def read_output_paths(cls, config) -> Dict[str, OutputPath]:
             }
         else:
             return None
-
-    @classmethod
-    def rename(cls, filepath, new_name: str) -> ExptConfig:
-        with open(filepath, "r") as f:
-            config = yaml.safe_load(f)
-            config["name"] = new_name
-
-        with open(filepath, "w") as f:
-            yaml.dump(config, f, sort_keys=False)
-
-        return ExptConfig(
-            workspace_id=config["workspace_id"],
-            unique_id=config["unique_id"],
-            name=config["name"],
-            started_at=config.get("started_at"),
-            finished_at=config.get("finished_at"),
-            success=config.get("success", "running"),
-            hasNWB=config["hasNWB"],
-            function=cls.read_function(config["function"]),
-            nwb=config.get("nwb"),
-            snakemake=config.get("snakemake"),
-        )

studio/app/common/core/experiment/experiment_writer.py

@@ -1,8 +1,11 @@
 import os
+import shutil
 from dataclasses import asdict
 from datetime import datetime
 from typing import Dict
 
+import yaml
+
 from studio.app.common.core.experiment.experiment import ExptConfig, ExptFunction
 from studio.app.common.core.experiment.experiment_builder import ExptConfigBuilder
 from studio.app.common.core.experiment.experiment_reader import ExptConfigReader
@@ -98,3 +101,50 @@ def function_from_nodeDict(self) -> ExptConfig:
                 func_dict[node.id].success = "success"
 
         return self.builder.set_function(func_dict).build()
+
+
+class ExptDataWriter:
+    def __init__(
+        self,
+        workspace_id: str,
+        unique_id: str,
+    ):
+        self.workspace_id = workspace_id
+        self.unique_id = unique_id
+
+    def delete_data(self) -> bool:
+        shutil.rmtree(
+            join_filepath([DIRPATH.OUTPUT_DIR, self.workspace_id, self.unique_id])
+        )
+        return True
+
+    def rename(self, new_name: str) -> ExptConfig:
+        filepath = join_filepath(
+            [
+                DIRPATH.OUTPUT_DIR,
+                self.workspace_id,
+                self.unique_id,
+                DIRPATH.EXPERIMENT_YML,
+            ]
+        )
+
+        # Note: "r+" option is not used here because it requires file pointer control.
+        with open(filepath, "r") as f:
+            config = yaml.safe_load(f)
+            config["name"] = new_name
+
+        with open(filepath, "w") as f:
+            yaml.dump(config, f, sort_keys=False)
+
+        return ExptConfig(
+            workspace_id=config["workspace_id"],
+            unique_id=config["unique_id"],
+            name=config["name"],
+            started_at=config.get("started_at"),
+            finished_at=config.get("finished_at"),
+            success=config.get("success", "running"),
+            hasNWB=config["hasNWB"],
+            function=ExptConfigReader.read_function(config["function"]),
+            nwb=config.get("nwb"),
+            snakemake=config.get("snakemake"),
+        )

studio/app/common/routers/experiment.py

@@ -1,5 +1,4 @@
 import os
-import shutil
 from glob import glob
 from typing import Dict
 
@@ -8,6 +7,7 @@
 
 from studio.app.common.core.experiment.experiment import ExptConfig
 from studio.app.common.core.experiment.experiment_reader import ExptConfigReader
+from studio.app.common.core.experiment.experiment_writer import ExptDataWriter
 from studio.app.common.core.utils.filepath_creater import join_filepath
 from studio.app.common.core.workspace.workspace_dependencies import (
     is_workspace_available,
@@ -45,12 +45,10 @@ async def get_experiments(workspace_id: str):
     dependencies=[Depends(is_workspace_owner)],
 )
 async def rename_experiment(workspace_id: str, unique_id: str, item: RenameItem):
-    config = ExptConfigReader.rename(
-        join_filepath(
-            [DIRPATH.OUTPUT_DIR, workspace_id, unique_id, DIRPATH.EXPERIMENT_YML]
-        ),
-        new_name=item.new_name,
-    )
+    config = ExptDataWriter(
+        workspace_id,
+        unique_id,
+    ).rename(item.new_name)
     config.nodeDict = []
     config.edgeDict = []
 
@@ -64,7 +62,10 @@ async def rename_experiment(workspace_id: str, unique_id: str, item: RenameItem)
 )
 async def delete_experiment(workspace_id: str, unique_id: str):
     try:
-        shutil.rmtree(join_filepath([DIRPATH.OUTPUT_DIR, workspace_id, unique_id]))
+        ExptDataWriter(
+            workspace_id,
+            unique_id,
+        ).delete_data()
         return True
     except Exception:
         return False
@@ -77,10 +78,11 @@ async def delete_experiment(workspace_id: str, unique_id: str):
 )
 async def delete_experiment_list(workspace_id: str, deleteItem: DeleteItem):
     try:
-        [
-            shutil.rmtree(join_filepath([DIRPATH.OUTPUT_DIR, workspace_id, uid]))
-            for uid in deleteItem.uidList
-        ]
+        for unique_id in deleteItem.uidList:
+            ExptDataWriter(
+                workspace_id,
+                unique_id,
+            ).delete_data()
         return True
     except Exception:
         return False