Belle II Software light-2406-ragdoll
pickable_basf2.py
1#!/usr/bin/env python3
2
3
10
11"""
12Pickable basf2
13Import this module at the top of your steering file to get a (more) pickable version of basf2.
14This is useful if you want to pickle the path using --dump-path and execute it later using --execute-path.
15Usually all calls to functions in basf2 like use_central_database are forgotten if you only save the path.
16With this module, these functions are executed again if you execute the pickled path using --execute-path.
17
18Technically this works by recording all calls to basf2 functions.
19If you have other code which should be pickled as well you can wrap it in the make_code_pickable function.
20
21If you want to exclude some functions, delete them from this module using
22del pickable_basf2.functionname
23"""
24
25
26import sys
27import pickle
28import inspect
29import unittest.mock as mock
30sys.modules['original_basf2'] = sys.modules['basf2'] # noqa
31import original_basf2
32
33
35 """ Drop-in replacement of the basf2 module, which keeps track of all functions calls """
36
37 def __getattr__(self, name):
38 """ Return attribute with the given name in the basf2 namespace """
39 return getattr(original_basf2, name)
40
41
42
43basf2_state_recorder = BASF2StateRecorder()
44
45manager = mock.Mock()
46
47
48def process(path, max_event=0):
49 """ Process call which pickles the recorded state in addition to the path """
50 sys.modules['basf2'] = original_basf2
51 original_basf2.process(path, max_event)
52 state = list(map(tuple, manager.mock_calls))
53 pickle_path = original_basf2.get_pickle_path()
54 print("Path", path)
55 print("State", state)
56 if pickle_path != '' and path is not None:
57 serialized = original_basf2.serialize_path(path)
58 serialized['state'] = state
59 pickle.dump(serialized, open(pickle_path, 'bw'))
60 sys.modules['basf2'] = basf2_state_recorder
61
62
63for name, x in original_basf2.__dict__.items():
64 # We record function and fake Boost.Python.function objects
65 if inspect.isfunction(x) or isinstance(x, type(original_basf2.find_file)):
66
67 mock_x = mock.Mock(x, side_effect=x)
68 manager.attach_mock(mock_x, name)
69 setattr(basf2_state_recorder, name, mock_x)
70 # Other names we have to set as well, because otherwise they won't be important by from basf2 import *
71 else:
72 setattr(basf2_state_recorder, name, x)
73
74
75basf2_state_recorder.process = process
76# and replace the module
77sys.modules['basf2'] = basf2_state_recorder
78sys.modules['pickable_basf2'] = basf2_state_recorder