Coverage for src/shephex/experiment/procedure/script.py: 95%
61 statements
« prev ^ index » next coverage.py v7.6.1, created at 2025-06-20 14:13 +0200
« prev ^ index » next coverage.py v7.6.1, created at 2025-06-20 14:13 +0200
1import importlib.util
2import inspect
3from pathlib import Path
4from typing import Callable, Optional, Union
6from shephex.experiment.context import ExperimentContext
7from shephex.experiment.options import Options
8from shephex.experiment.result import ExperimentResult
10from .pickle import PickleProcedure
11from .procedure import Procedure
14class ScriptProcedure(Procedure):
15 """
16 A procedure wrapping a script.
17 """
19 def __init__(
20 self,
21 function: Optional[Callable] = None,
22 path: Optional[Path | str] = None,
23 function_name: Optional[str] = None,
24 code: Optional[str] = None,
25 context: bool = False,
26 ) -> None:
27 super().__init__(name='procedure.py', context=context)
28 if function is not None:
29 self.function_name = function.__name__
30 self.script_code = Path(inspect.getsourcefile(function)).read_text()
31 elif path is not None and function_name is not None:
32 if not isinstance(path, Path):
33 path = Path(path)
34 self.function_name = function_name
35 self.script_code = path.read_text()
36 elif code is not None and function_name is not None:
37 self.function_name = function_name
38 self.script_code = code
39 else:
40 raise ValueError(
41 'ScriptProcedure requires one of following sets of arguments: function, path and function_name, or code and function_name.'
42 )
44 def dump(self, directory: Union[Path, str]) -> None:
45 directory = Path(directory)
46 with open(directory / self.name, 'w') as f:
47 f.write(self.script_code)
49 def _execute(
50 self,
51 options: Options,
52 directory: Optional[Union[Path, str]] = None,
53 shephex_directory: Optional[Union[Path, str]] = None,
54 context: Optional[ExperimentContext] = None,
55 ) -> ExperimentResult:
56 """
57 Execute the procedure by running the script on a subprocess.
59 Parameters
60 ----------
61 directory : Optional[Union[Path, str]], optional
62 """
63 # Directory handling
64 if directory is None:
65 directory = Path.cwd() # pragma: no cover
67 # Basic command
68 path = (shephex_directory / self.name).resolve()
70 assert (
71 shephex_directory.exists()
72 ), f'Directory {shephex_directory} does not exist.'
73 assert path.exists(), f'File {path} does not exist.'
74 assert Path(directory).exists(), f'Directory {directory} does not exist.'
76 func_function = self.get_function_from_script(path, self.function_name)
77 if hasattr(func_function, '__wrapped__'): # check if decorated
78 func = func_function()
79 else:
80 func = func_function
82 func_procedure = PickleProcedure(func, context=self.context)
84 return func_procedure._execute(options, directory, shephex_directory, context)
86 def hash(self) -> int:
87 return self.script_code.__hash__()
89 def get_function_from_script(self, script_path: Path, func_name: str) -> None:
90 """Dynamically imports a function from a Python script given its path."""
91 script_path = Path(script_path).resolve() # Ensure absolute path
92 module_name = script_path.stem # Extract filename without .py
94 # Create a module spec
95 spec = importlib.util.spec_from_file_location(module_name, script_path)
96 if spec is None:
97 raise ImportError(f'Could not load spec for {script_path}')
99 # Create a module from the spec
100 module = importlib.util.module_from_spec(spec)
101 spec.loader.exec_module(module) # Execute the module
103 # Retrieve the function
104 if not hasattr(module, func_name):
105 raise AttributeError(f"Module {module_name} has no function '{func_name}'")
107 return getattr(module, func_name)
109 def get_metadata(self) -> dict:
110 metadata = super().get_metadata()
111 metadata['type'] = 'ScriptProcedure'
112 metadata['function_name'] = self.function_name
113 return metadata
115 @classmethod
116 def from_metadata(cls, metadata: dict) -> 'ScriptProcedure':
117 return cls(
118 function_name=metadata['function_name'],
119 context=metadata['context'],
120 path=metadata['path'],
121 )