cleaned code for remove step cache for local workspaces

tango/step_caches/local_step_cache.py

@@ -2,6 +2,7 @@
 import logging
 import warnings
 import weakref
+import shutil
 from pathlib import Path
 from typing import Any, MutableMapping, Optional, OrderedDict, Union, cast
 
@@ -147,6 +148,13 @@ def __setitem__(self, step: Step, value: Any) -> None:
                 pass
             raise
 
+    def __delitem__(self, step_unique_id) -> None:
+        location = self.dir / step_unique_id
+        try:
+            shutil.rmtree(location)
+        except OSError:
+            raise OSError('Step Cache folder not found')
+
     def __len__(self) -> int:
         return sum(1 for _ in self.dir.glob(f"*/{self.METADATA_FILE_NAME}"))
 

tango/workspace.py

@@ -422,27 +422,14 @@ def step_result(self, step_name: str) -> Any:
         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 step_cache_remove(self, step_unique_id: str) -> Any:
+        """
+        Removes cached step using the given unique step id
+        :raises KeyError: If there is no step with the given name.
+        """
+        raise NotImplementedError()
+
+
 
     def capture_logs_for_run(self, name: str) -> ContextManager[None]:
         """

tango/workspaces/local_workspace.py

@@ -322,6 +322,19 @@ def step_failed(self, step: Step, e: BaseException) -> None:
             lock.release()
             del self.locks[step]
 
+    def step_cache_remove(self, step_unique_id: str) -> None:
+        """
+        Get Step 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
+        """
+        with SqliteDict(self.step_info_file) as d:
+            try:
+                del d[step_unique_id]
+                d.commit()
+                self.cache.__delitem__(step_unique_id)
+            except KeyError:
+                raise KeyError(f"No step named '{step_unique_id}' found")
+
     def register_run(self, targets: Iterable[Step], name: Optional[str] = None) -> Run:
         # sanity check targets
         targets = list(targets)