[experiment](feat) Refactor as context manager

Author: kzajac97

pydentification/experiment/storage/compose.py

@@ -1,10 +1,39 @@
 import importlib.util
 import json
 import os
+import shutil
 import sys
 import zipfile
 from pathlib import Path
-from typing import Callable
+from typing import Any, Callable
+
+
+class ReplaceSourceCode:
+    """
+    ContextManager over-writing imports with given `path` to ZIP with source code created by
+    `pydentification.experiment.storage.code.save_code_snapshot`.
+
+    Code is extracted to a temporary directory and added to the `sys.path` for the duration of the context and removed
+    afterward on exit. The source code needs to be unique directory to avoid conflicts with other imports.
+    """
+
+    def __init__(self, path: Path):
+        self.path = path
+        self.source_path = path.with_suffix("")
+
+    def __enter__(self):
+        if self.source_path.exists():
+            raise FileExistsError(f"Can't overwrite {self.source_path.stem}!")
+
+        with zipfile.ZipFile(self.path, "r") as zip:
+            zip.extractall(str(self.source_path))
+
+        sys.path.append(str(self.source_path))
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        sys.path.remove(str(self.source_path))
+        shutil.rmtree(self.source_path)
 
 
 def _import_function_from_path(module_path: str, function_name: str) -> Callable:
@@ -21,12 +50,9 @@ def _import_function_from_path(module_path: str, function_name: str) -> Callable
     return function
 
 
-def _safe_unzip(path: Path):
-    if path.with_suffix("").exists():
-        raise FileExistsError(f"Can't overwrite {path.stem}!")
-
-    with zipfile.ZipFile(path, "r") as zip:
-        zip.extractall(str(path.parent))
+def _load_model_and_parameters(path: str | Path, name: str, parameters: dict[str, Any]) -> Any:
+    model_fn = _import_function_from_path(path, name)
+    return model_fn(**parameters)
 
 
 def compose_model(
@@ -48,20 +74,14 @@ def compose_model(
     if isinstance(source, str):
         source = Path(source)
 
-    if source is not None:
-        _safe_unzip(source)
-        sys.path.append(str(source.parent))
-
-    model_fn = _import_function_from_path(path, name)
-
     if parameters is not None:
         with open(parameters, "r") as f:
             parameters = json.load(f)
     else:
         parameters = {}
 
-    model = model_fn(**parameters)
-
     if source is not None:
-        sys.path.remove(str(source.parent))
-    return model
+        with ReplaceSourceCode(source):
+            return _load_model_and_parameters(path, name, parameters)
+    else:
+        return _load_model_and_parameters(path, name, parameters)