Refactor context import and usage in capsula module

capsula/_decorator.py

@@ -9,7 +9,7 @@
 
 from ._backport import ParamSpec
 from ._context import ContextBase
-from ._run import CapsuleParams, Run
+from ._run import CapsuleParams, FuncInfo, Run
 from ._watcher import WatcherBase
 
 if TYPE_CHECKING:
@@ -108,7 +108,7 @@ def decorator(func_or_run: Callable[_P, _T] | Run[_P, _T]) -> Run[_P, _T]:
 
 def context(
     context: ContextBase | Callable[[CapsuleParams], ContextBase],
-    mode: Literal["pre", "in", "post", "all"],
+    mode: Literal["pre", "post", "all"],
 ) -> Callable[[Callable[_P, _T] | Run[_P, _T]], Run[_P, _T]]:
     def decorator(func_or_run: Callable[_P, _T] | Run[_P, _T]) -> Run[_P, _T]:
         func = func_or_run.func if isinstance(func_or_run, Run) else func_or_run
@@ -120,7 +120,7 @@ def decorator(func_or_run: Callable[_P, _T] | Run[_P, _T]) -> Run[_P, _T]:
 
 
 def run(
-    run_dir: Path | Callable[[CapsuleParams], Path],
+    run_dir: Path | Callable[[FuncInfo], Path],
 ) -> Callable[[Callable[_P, _T] | Run[_P, _T]], Run[_P, _T]]:
     def decorator(func_or_run: Callable[_P, _T] | Run[_P, _T]) -> Run[_P, _T]:
         func = func_or_run.func if isinstance(func_or_run, Run) else func_or_run

capsula/capture.py

@@ -5,13 +5,13 @@
 import logging
 import subprocess
 
-from capsula._context import ContextBase
 from capsula.config import CapsulaConfig
+from capsula.context import Context
 
 logger = logging.getLogger(__name__)
 
 
-def capture(*, config: CapsulaConfig) -> ContextBase:
+def capture(*, config: CapsulaConfig) -> Context:
     """Capture the context."""
     logger.debug(f"Capture config: {config.capture}")
 
@@ -36,7 +36,7 @@ def capture(*, config: CapsulaConfig) -> ContextBase:
 
     config.ensure_capsule_directory_exists()
 
-    ctx = ContextBase.capture(config)
+    ctx = Context.capture(config)
 
     # Write the context to the output file.
     with (config.capsule / "context.json").open("w") as output_file: