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

1import importlib.util 

2import inspect 

3from pathlib import Path 

4from typing import Callable, Optional, Union 

5 

6from shephex.experiment.context import ExperimentContext 

7from shephex.experiment.options import Options 

8from shephex.experiment.result import ExperimentResult 

9 

10from .pickle import PickleProcedure 

11from .procedure import Procedure 

12 

13 

14class ScriptProcedure(Procedure): 

15 """ 

16 A procedure wrapping a script. 

17 """ 

18 

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 ) 

43 

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) 

48 

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. 

58 

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 

66 

67 # Basic command 

68 path = (shephex_directory / self.name).resolve() 

69 

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.' 

75 

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 

81 

82 func_procedure = PickleProcedure(func, context=self.context) 

83 

84 return func_procedure._execute(options, directory, shephex_directory, context) 

85 

86 def hash(self) -> int: 

87 return self.script_code.__hash__() 

88 

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 

93 

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}') 

98 

99 # Create a module from the spec 

100 module = importlib.util.module_from_spec(spec) 

101 spec.loader.exec_module(module) # Execute the module 

102 

103 # Retrieve the function 

104 if not hasattr(module, func_name): 

105 raise AttributeError(f"Module {module_name} has no function '{func_name}'") 

106 

107 return getattr(module, func_name) 

108 

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 

114 

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 )