Merge pull request #59 from cyber-physical-systems-group/stubs/path-casting

Stubs/path casting

Author: kzajac97

pydentification/experiment/storage/code.py

@@ -3,6 +3,8 @@
 import uuid
 from pathlib import Path
 
+from pydentification.stubs import cast_to_path
+
 
 def _load_gitignore() -> set[str]:
     """Load .gitignore from default name and root directory as set"""
@@ -31,6 +33,7 @@ def _skip_subdir(current: Path, archive_path: Path, forbidden_paths: set[str]) -
     return False
 
 
+@cast_to_path
 def save_code_snapshot(
     name: str,
     source_dir: str | Path,
@@ -49,12 +52,6 @@ def save_code_snapshot(
     :param accept_suffix: set of suffixes to include in the archive
     :param use_gitignore: whether to use .gitignore file in the source directory for filter_prefix
     """
-    if isinstance(source_dir, str):
-        source_dir = Path(source_dir)
-
-    if isinstance(target_dir, str):
-        target_dir = Path(target_dir)
-
     source_dir = Path(source_dir).resolve()  # ensure absolute path
     snapshot_path = target_dir / name
     temp_dir = target_dir / str(uuid.uuid4())  # create temp dir with unique name for copying files

pydentification/experiment/storage/compose.py

@@ -7,6 +7,8 @@
 from pathlib import Path
 from typing import Any, Callable
 
+from pydentification.stubs import cast_to_path
+
 
 class ReplaceSourceCode:
     """
@@ -55,6 +57,7 @@ def _load_model_and_parameters(path: str | Path, name: str, parameters: dict[str
     return model_fn(**parameters)
 
 
+@cast_to_path
 def compose_model(
     path: str | Path,
     name: str = "model_fn",
@@ -71,9 +74,6 @@ def compose_model(
     :param source: filesystem Path to the ZIP file with source code
                    if None imports are attempted from the current working directory.
     """
-    if isinstance(source, str):
-        source = Path(source)
-
     if parameters is not None:
         with open(parameters, "r") as f:
             parameters = json.load(f)

pydentification/experiment/storage/models.py

@@ -6,6 +6,8 @@
 import torch
 from safetensors.torch import save_model
 
+from pydentification.stubs import cast_to_path
+
 
 def save_torch(path: Path, model: torch.nn.Module, method: Literal["pt", "safetensors"] = "safetensors"):
     if method == "safetensors":
@@ -21,6 +23,7 @@ def save_json(path: Path, data: dict):
         json.dump(data, f)  # type: ignore
 
 
+@cast_to_path
 def save_lightning(
     path: str | Path,
     model: pl.LightningModule,
@@ -33,10 +36,7 @@ def save_lightning(
     :param method: method of saving the model, either "pt" or "safetensors"
     :param save_hparams: whether to save hyperparameters in a JSON file
     """
-    if isinstance(path, str):
-        path = Path(path)
     path.mkdir(parents=True, exist_ok=True)
-
     save_torch(path / f"trained-model.{method}", model=model.module, method=method)  # save only the model
 
     if save_hparams:

pydentification/experiment/storage/sync.py

@@ -2,11 +2,11 @@
 
 import wandb
 
+from pydentification.stubs import cast_to_path
 
+
+@cast_to_path
 def save_to_wandb(path: str | Path):
     """Save all files from directory to W&B"""
-    if isinstance(path, str):
-        path = Path(path)
-
     for file in path.rglob("*"):
         wandb.save(str(file))  # save file one by one

pydentification/stubs.py

@@ -1,3 +1,37 @@
-from typing import Callable
+from functools import wraps
+from pathlib import Path
+from typing import Any, Callable, Union, get_type_hints
 
 Print = Callable[[str], None]
+
+
+def cast_to_path(func: Callable):
+    """
+    Decorator to cast string arguments to Path in the function signature.
+    Uses type hints to determine which arguments should be cast.
+
+    :note: do not use with functions that have missing type hints
+    """
+
+    def cast(arg: Any, hint: Any) -> Any:
+        if hint == Union[str, Path]:
+            return Path(arg)
+        return arg
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        hints = get_type_hints(func)
+        new_args = []
+        new_kwargs = {}
+
+        if len(args) + len(kwargs) > len(hints):
+            raise TypeError("Some arguments are missing type hints!")
+
+        for arg, (name, hint) in zip(args, hints.items()):
+            new_args.append(cast(arg, hint))
+        for name, arg in kwargs.items():
+            new_kwargs[name] = cast(arg, hint=hints.get(name))
+
+        return func(*new_args, **new_kwargs)
+
+    return wrapper