From 173a2d34ed6f77cddde682b50c0bff6977acd703 Mon Sep 17 00:00:00 2001
From: itutu-tienday <>
Date: Thu, 21 Nov 2024 16:09:58 +0900
Subject: [PATCH] Add ExptDataWriter - porting from
 arayabrain/barebone-studio#366

---
 .../core/experiment/experiment_reader.py      | 22 --------
 .../core/experiment/experiment_writer.py      | 50 +++++++++++++++++++
 studio/app/common/routers/experiment.py       | 26 +++++-----
 3 files changed, 64 insertions(+), 34 deletions(-)

diff --git a/studio/app/common/core/experiment/experiment_reader.py b/studio/app/common/core/experiment/experiment_reader.py
index d092a77fb..6789d4ba4 100644
--- a/studio/app/common/core/experiment/experiment_reader.py
+++ b/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"),
-        )
diff --git a/studio/app/common/core/experiment/experiment_writer.py b/studio/app/common/core/experiment/experiment_writer.py
index 4d49a1c67..7e02b0fd5 100644
--- a/studio/app/common/core/experiment/experiment_writer.py
+++ b/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"),
+        )
diff --git a/studio/app/common/routers/experiment.py b/studio/app/common/routers/experiment.py
index c8b8a6738..04d4e9518 100644
--- a/studio/app/common/routers/experiment.py
+++ b/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