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

1import importlib.util 

2import inspect 

3import types 

4from pathlib import Path 

5from typing import Callable, Optional, Union 

6 

7from shephex.experiment.context import ExperimentContext 

8from shephex.experiment.options import Options 

9from shephex.experiment.result import ExperimentResult 

10 

11from .pickle import PickleProcedure 

12from .procedure import Procedure 

13 

14 

15class ScriptProcedure(Procedure): 

16 """ 

17 A procedure wrapping a script. 

18 """ 

19 

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 ) 

44 

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) 

49 

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. 

59 

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 

67 

68 # Basic command 

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

70 

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

76 

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 

82 

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

84 

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

86 

87 def hash(self) -> int: 

88 return self.script_code.__hash__() 

89 

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 

94 

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

99 

100 # Create a module from the spec 

101 module = importlib.util.module_from_spec(spec) 

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

103 

104 # Retrieve the function 

105 if not hasattr(module, func_name): 

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

107 

108 return getattr(module, func_name) 

109 

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 

115 

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 )