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