Belle II Software light-2406-ragdoll
pickle_path.py
1#!/usr/bin/env python3
2
3
10
11"""
12basf2.pickle_path - Functions necessary to pickle and unpickle a Path
13=====================================================================
14
15This module contains all the functiones necessary to serialize and deserialize
16a full path with all modules, parameters, sub paths, conditions and so on. This
17can be used in conjunction with ``basf2 --dump-path`` and ``basf2
18--execute-path`` to save a full configuration to file and execute it later.
19"""
20
21import pybasf2
22import pickle as _pickle
23import os as _os
24import sys as _sys
25
26
27def serialize_value(module, parameter):
28 """Serialize a single basf2 module parameter"""
29 if parameter.name == 'path' and module.type() == 'SubEvent':
30 return serialize_path(parameter.values)
31 else:
32 return parameter.values
33
34
35def deserialize_value(module, parameter_state):
36 """Deserialize a single basf2 module paramater"""
37 if parameter_state['name'] == 'path' and module.type() == 'SubEvent':
38 return deserialize_path(parameter_state['values'])
39 else:
40 return parameter_state['values']
41
42
43def serialize_conditions(module):
44 """Serialize all conditions attached to a basf2 module"""
45 condition_list = []
46
47 for condition in module.get_all_conditions():
48 condition_list.append({'value': condition.get_value(),
49 'operator': int(condition.get_operator()),
50 'path': serialize_path(condition.get_path()),
51 'option': int(condition.get_after_path())})
52
53 return condition_list
54
55
56def deserialize_conditions(module, module_state):
57 """Deserialize all conditions for a given basf2 module"""
58 conditions = module_state['condition']
59 for cond in conditions:
60 module.if_value(str(pybasf2.ConditionOperator.values[cond['operator']]) + str(cond['value']),
61 deserialize_path(cond['path']), pybasf2.AfterConditionPath.values[cond['option']])
62
63
64def serialize_module(module):
65 """Serialize a basf2 module into a python dictionary. Doesn't work for python modules"""
66 if module.type() == '' or module.type() == 'PyModule':
67 raise RuntimeError(
68 f"Module '{module.name()}' doesn't have a type or is a Python module! Note that --dump-path cannot work properly " +
69 "with basf2 modules written in Python.")
70 return {
71 'name': module.name(),
72 'type': module.type(),
73 'flag': module.has_properties(pybasf2.ModulePropFlags.PARALLELPROCESSINGCERTIFIED),
74 'parameters': [{'name': parameter.name, 'values': serialize_value(module, parameter)}
75 for parameter in module.available_params()
76 if parameter.setInSteering or module.type() == 'SubEvent'],
77 'condition': serialize_conditions(module) if module.has_condition() else None}
78
79
80def deserialize_module(module_state):
81 """Deserialize a basf2 module from a python dictionary"""
82 module = pybasf2._register_module(module_state['type'])
83 module.set_name(module_state['name'])
84 if 'condition' in module_state and module_state['condition'] is not None:
85 deserialize_conditions(module, module_state)
86 if 'flag' in module_state and module_state['flag']:
87 # for some modules, this flag might be changed from the default
88 module.set_property_flags(pybasf2.ModulePropFlags.PARALLELPROCESSINGCERTIFIED)
89 for parameter_state in module_state['parameters']:
90 module.param(parameter_state['name'],
91 deserialize_value(module, parameter_state))
92 return module
93
94
95def serialize_path(path):
96 """Serialize a basf2 Path into a python dictionary"""
97 return {'modules': [serialize_module(module) for module in path.modules()]}
98
99
100def deserialize_path(path_state):
101 """Deserialize a basf2 Path from a python dictionary"""
102 path = pybasf2.Path()
103 for module_state in path_state['modules']:
104 module = deserialize_module(module_state)
105 path.add_module(module)
106 return path
107
108
109def get_path_from_file(path_filename):
110 """Read a path from a given pickle file"""
111 with open(path_filename, 'br') as f:
112 return deserialize_path(_pickle.load(f))
113
114
115def write_path_to_file(path, filename):
116 """Write a path to a given pickle file"""
117 with open(filename, 'bw') as f:
118 _pickle.dump(serialize_path(path), f)
119
120
121def check_pickle_path(path):
122 """Check if the path to be executed should be pickled or unpickled.
123 This function is used by basf2.process to handle the ``--dump-path`` and
124 ``--execute-path`` arguments to ``basf2``
125 """
126 # If a pickle path is set via --dump-path or --execute-path we do something special
127 pickle_filename = pybasf2.get_pickle_path()
128 if pickle_filename == '':
129 return path
130
131 # If the given path is None and the picklePath is valid we load a path from the pickle file
132 if _os.path.isfile(pickle_filename) and path is None:
133 path = get_path_from_file(pickle_filename)
134 with open(pickle_filename, "br") as f:
135 loaded = _pickle.load(f)
136 if 'state' in loaded:
137 pybasf2.B2INFO("Pickled path contains a state object. Activating pickled state.")
138 for name, args, kwargs in loaded['state']:
139 getattr(_sys.modules[__name__], name)(*args, **kwargs)
140 return path
141
142 # Otherwise we dump the given path into the pickle file and exit
143 elif path is not None:
144 write_path_to_file(path, pickle_filename)
145 return None
146 else:
147 pybasf2.B2FATAL("Couldn't open path-file '" + pickle_filename + "' and no steering file provided.")
148
149
150def make_code_pickable(code):
151 """
152 Sometimes it is necessary to execute code which won't be pickled if a user dumps the basf2 path and wants to execute it later. Using the pickable_basf2 module all calls to basf2 functions
153 are recorded. Now if a user has to execute code outside of basf2, e.g. modifying objects in the ROOT namespace,
154 this won't be pickled. By wrapping the code in this function it is technically a call to a basf2 function
155 and will be pickled again. Problem solved.
156 """
157 exec(code, globals())
158