From ab727838e1be5c96c071048fca96738a0aa784b1 Mon Sep 17 00:00:00 2001
From: Pranjali Basmatkar <pranjalib@Pranjali-Basmatkars-MacBook-Pro.local>
Date: Wed, 12 Jul 2023 11:42:39 -0700
Subject: [PATCH] remove step feature for local workspaces

---
 tango/workspace.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/tango/workspace.py b/tango/workspace.py
index e9bb7281..a7c671a7 100644
--- a/tango/workspace.py
+++ b/tango/workspace.py
@@ -20,6 +20,8 @@
 from urllib.parse import ParseResult, urlparse
 
 import pytz
+from sqlitedict import SqliteDict
+import shutil
 
 from .common import Registrable
 from .common.from_params import FromParams
@@ -419,6 +421,29 @@ def step_result(self, step_name: str) -> Any:
                 return self.step_cache[run.steps[step_name]]
         raise KeyError(f"No step named '{step_name}' found in previous runs")
 
+
+    def remove_step(self, step_name: str) -> Any:
+        """
+        Get Step name (Unique ID) from the user and remove the step information from cache
+        :raises KeyError: If no step with the unique name found in the cache dir
+        """
+        # get path to dir dynamically
+        sqlite_path = self.dir / "stepinfo.sqlite"
+        with SqliteDict(sqlite_path) as d:
+            try:
+                step_location = d[step_name].result_location
+                # remove step info from the sqlite dict
+                del d[step_name]
+                d.commit()
+                # remove cache directory
+                try:
+                    shutil.rmtree(step_location)
+                except OSError:
+                    raise OSError('Step Cache folder not found')
+                return('Step deleted')
+            except KeyError:
+                raise KeyError(f"No step named '{step_name}' found")
+
     def capture_logs_for_run(self, name: str) -> ContextManager[None]:
         """
         Should return a context manager that can be used to capture the logs for a run.