Coverage for src/shephex/experiment/procedure/pickle.py: 98%

61 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2025-06-20 14:13 +0200

1import os 

2import pickle 

3import traceback 

4from inspect import getsource, getsourcefile, signature 

5from pathlib import Path 

6from typing import Callable, Optional, Union 

7 

8import dill 

9 

10from shephex.experiment.context import ExperimentContext 

11from shephex.experiment.options import Options 

12from shephex.experiment.result import ExperimentError, ExperimentResult 

13from shephex.experiment.status import Status 

14 

15from .procedure import Procedure 

16 

17 

18class PickleProcedure(Procedure): 

19 """ 

20 A procedure wrapping a function. 

21 """ 

22 def __init__(self, func: Callable, context: bool = False) -> None: 

23 super().__init__(name='procedure.pkl', context=context) 

24 self.func = func 

25 self.signature = signature(func) 

26 self.script_code = Path(getsourcefile(func)).read_text() 

27 

28 

29 def dump(self, directory: Union[Path, str]) -> None: 

30 directory = Path(directory) 

31 with open(directory / self.name, 'wb') as f: 

32 dill.dump(self, f, protocol=pickle.HIGHEST_PROTOCOL, recurse=True) 

33 

34 name = self.name.replace('.pkl', '.py') 

35 with open(directory / name, 'w') as f: 

36 f.write(self.script_code) 

37 

38 def _execute( 

39 self, 

40 options: Options, 

41 directory: Optional[Union[Path, str]] = None, 

42 shephex_directory: Optional[Union[Path, str]] = None, 

43 context: Optional[ExperimentContext] = None, 

44 ) -> ExperimentResult: 

45 # Directory handling 

46 """ 

47 Execute the procedure by calling the function. 

48 

49 Parameters 

50 ---------- 

51 options : Options 

52 The options for the procedure. 

53 directory : Optional[Union[Path, str]], optional 

54 Directory where the procedure will be executed, by default None. 

55 shephex_directory : Optional[Union[Path, str]], optional 

56 Directory where the procedure is saved, by default None. 

57 context : Optional[ExperimentContext], optional 

58 The context for the experiment, by default None. 

59 

60 Returns 

61 ------- 

62 ExperimentResult 

63 The result of the experiment. 

64 """ 

65 

66 cwd = Path.cwd() 

67 if directory is not None: 

68 directory = Path(directory) 

69 directory.mkdir(parents=True, exist_ok=True) 

70 os.chdir(directory) 

71 

72 if context is None: # Execution without context 

73 try: 

74 result = self.func(*options.args, **options.kwargs) 

75 status = Status.completed() 

76 except Exception as e: 

77 print(traceback.format_exc()) 

78 result = ExperimentError(e) 

79 status = Status.failed() 

80 else: # Execution with context 

81 if 'context' in options.kwargs.keys(): 

82 raise ValueError( 

83 '"context" is a reserved keyword if shephex context is enabled.' 

84 ) 

85 

86 try: 

87 result = self.func(*options.args, **options.kwargs, context=context) 

88 status = Status.completed() 

89 

90 except Exception as e: 

91 print(traceback.format_exc()) 

92 result = ExperimentError(e) 

93 status = Status.failed() 

94 

95 result = ExperimentResult(result=result, status=status) 

96 result.dump(shephex_directory) 

97 # Directory handling 

98 os.chdir(cwd) 

99 return result 

100 

101 def hash(self) -> int: 

102 return hash(getsource(self.func)) 

103 

104 def get_metadata(self) -> dict: 

105 metadata = super().get_metadata() 

106 metadata['type'] = 'PickleProcedure' 

107 return metadata 

108 

109 @classmethod 

110 def from_metadata(cls, metadata: dict): 

111 raise NotImplementedError('PickleProcedure cannot be created from metadata - currently.')