Coverage for src/shephex/experiment/procedure/pickle.py: 98%
61 statements
« prev ^ index » next coverage.py v7.6.1, created at 2025-03-30 12:30 +0200
« prev ^ index » next coverage.py v7.6.1, created at 2025-03-30 12:30 +0200
1import os
2import pickle
3import traceback
4from inspect import getsource, getsourcefile, signature
5from pathlib import Path
6from typing import Callable, Optional, Union
8import dill
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
15from .procedure import Procedure
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()
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)
34 name = self.name.replace('.pkl', '.py')
35 with open(directory / name, 'w') as f:
36 f.write(self.script_code)
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.
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.
60 Returns
61 -------
62 ExperimentResult
63 The result of the experiment.
64 """
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)
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 )
86 try:
87 result = self.func(*options.args, **options.kwargs, context=context)
88 status = Status.completed()
90 except Exception as e:
91 print(traceback.format_exc())
92 result = ExperimentError(e)
93 status = Status.failed()
95 result = ExperimentResult(result=result, status=status)
96 result.dump(shephex_directory)
97 # Directory handling
98 os.chdir(cwd)
99 return result
101 def hash(self) -> int:
102 return hash(getsource(self.func))
104 def get_metadata(self) -> dict:
105 metadata = super().get_metadata()
106 metadata['type'] = 'PickleProcedure'
107 return metadata
109 @classmethod
110 def from_metadata(cls, metadata: dict):
111 raise NotImplementedError('PickleProcedure cannot be created from metadata - currently.')