Belle II Software light-2406-ragdoll
datastoreprinter.py
1#!/usr/bin/env python3
2
3
10
11from ROOT import Belle2, kIsPublic, kIsStatic, TVector3, TLorentzVector
12from ROOT.Belle2 import Const
13from ROOT.Math import XYZVector
14from basf2 import Module, B2FATAL
15
16
17def get_public_members(classname):
18 """
19 Return a list of public, non-static member functions for a given classname.
20 The class must exist in the Belle2 namespace and have a ROOT dictionary """
21 if not hasattr(Belle2, classname):
22 return []
23 tclass = getattr(Belle2, classname).Class()
24 members = {e.GetName() for e in tclass.GetListOfMethods()
25 if (e.Property() & kIsPublic) and not (e.Property() & kIsStatic)}
26
27 # filter known members from ClassDef and constructor/destructor
28 # Remove some Dictionary members
29 removed = {
30 "CheckTObjectHashConsistency", "Class", "Class_Name", "Class_Version",
31 "DeclFileLine", "DeclFileName", "Dictionary", "ImplFileLine",
32 "ImplFileName", "IsA", "ShowMembers", "Streamer", "StreamerNVirtual",
33 "operator!=", "operator=", "operator==",
34 # we don't need constructor and destructor either
35 classname, f"~{classname}",
36 }
37 members -= removed
38 return list(sorted(members))
39
40
41class DataStorePrinter:
42 """
43 Class to print contents of a StoreObjPtr or StoreArray.
44
45 This class is inteded to print the contents of dataobjects to the standard
46 output to monitor changes to the contents among versions.
47
48 For example:
49
50 >>> printer = DataStorePrinter("MCParticle", ["getVertex"], {"hasStatus": [1, 2, 4]})
51 >>> printer.print()
52
53 will loop over all MCParticle instances in the MCParticles StoreArray and
54 print someting like ::
55
56 MCParticle#0
57 getVertex(): (0,0,0)
58 hasStatus(1): True
59 hasStatus(2): False
60 hasStatus(4): False
61
62 for each MCparticle
63 """
64
65 def __init__(self, name, simple, withArgument=None, array=True):
66 """
67 Initialize
68
69 Args:
70 name (str): class name of the DataStore object
71 simple (list): list of member names to print which do not need any additional
72 arguments
73 withArgument (dict or None): dictionary of member names and a list of
74 all argument combinations to call them.
75 array (bool): if True we print a StoreArray, otherwise a single StoreObjPtr
76 """
77
78 self.name = name
80 self.array = array
82 self.object_members = []
84 # add the simple members to the list of members to call with empty
85 # arguments
86 for member in simple:
87 self.object_members.append((member, [], None, None))
88
89 # and add the members with Argument
90 if withArgument:
91 for member, arguments in withArgument.items():
92 for args in arguments:
93 # convert args to a list
94 if not isinstance(args, list) or isinstance(args, tuple):
95 args = [args]
96 # and add (name,args,display,callback) tuple
97 self.object_members.append((member, args, None, None))
98
99 # sort them by member name to have fixed order: python sort is
100 # guaranteed to be stable so different calls to the same member will
101 # remain in same relative order
102 self.object_members.sort(key=lambda x: x[0])
103
104 def add_member(self, name, arguments=None, print_callback=None, display=None):
105 """
106 Add an additional member to be printed.
107
108 Args:
109 name (str): name of the member
110 arguments (list or callable): arguments to pass to the member when calling it
111 If this is a callable object then the function is called with
112 the object as first argument and the member name to be tested as
113 second argument. The function is supposed to return the list of
114 arguments to pass to the member when calling. Possible return
115 valus for the callable are:
116
117 * a `list` of arguments to be passed to the member. An empty
118 `list` means to call the member with no arguments.
119 * a `tuple` of `lists <list>` to call the member once for each
120 list of arguments in the tuple
121 * `None` or an empty tuple to not call the member at all
122 print_callback (function or None): if not None a function to print
123 the result of the member call. The function will be called with
124 the arguments (name, arguments, result) and should print the
125 result on stdout without any additional information.
126 display (str or None): display string to use when printing member call
127 info instead of function name + arguments. If it is not given
128 the default output will be ``{membername}({arguments}):``
129 followed by the result.
130 """
131 if arguments is None:
132 arguments = []
133 self.object_members.append((name, arguments, print_callback, display))
134 # return self for method chaining
135 return self
136
137 def print(self):
138 """Print all the objects currently existing"""
139 if self.array:
140 data = Belle2.PyStoreArray(self.name + "s")
141 if not data.isValid():
142 print(f"No data for {self.name}")
143 else:
144 for i, obj in enumerate(data):
145 self._printObj(obj, i)
146 else:
147 obj = Belle2.PyStoreObj(self.name)
148 if not obj.isValid():
149 print(f"No data for {self.name}")
150 else:
151 if obj:
152 self._printObj(obj.obj())
153
154 def print_untested(self):
155 """Print all the public member functions we will not test"""
156 members = get_public_members(self.name)
157 tested = {e[0] for e in self.object_members}
158 for member in members:
159 if member in tested:
160 continue
161 print(f"Untested method {self.name}::{member}")
162
163 def _printObj(self, obj, index=None):
164 """Print all defined members for each object with given index.
165 If we print a StoreObjPtr then index is None and this function is called
166 once. If we print StoreArrays this function is called once for each
167 entry in the array with index set to the position in the array
168 """
169 # print array index? If yes then add it to the output
170 if index is not None:
171 index = f"#{int(index)}"
172 else:
173 index = ""
174
175 print(f"{self.name}{index}:")
176
177 # loop over all defined member/argument combinations
178 # and print "member(arguments): result" for each
179 for name, arguments, callback, display in self.object_members:
180 # if arguments is callable it means we need to call it to determine
181 # the arguments
182 if callable(arguments):
183 all_args = arguments(obj, name)
184 # None means we don't calle the member this time
185 if all_args is None:
186 continue
187 # list is one set of arguments, tuple(list) is n set of
188 # arguments. So convert list into tuple of length 1 to call
189 # member once
190 if isinstance(all_args, list):
191 all_args = (all_args,)
192 else:
193 # if arguments is not callable we assume it's one set of
194 # arguments
195 all_args = (arguments,)
196
197 for args in all_args:
198 result = getattr(obj, name)(*args)
199 # display can be used to override what is printed for the member. If
200 # so, use it here
201 if display is not None:
202 print(" " + display + ": ", end="")
203 else:
204 # otherwise just print name and arguments
205 print(f" {name}({','.join(map(str, args))}): ", end="")
206 # if a callback is set the callback is used to print the result
207 if callback is not None:
208 print("", end="", flush=True)
209 callback(name, args, result)
210 # otherwise use default function
211 else:
212 self._printResult(result)
213
214 def _printResult(self, result, depth=0, weight=None):
215 """ Print the result of calling a certain member.
216 As we want the test to be independent of memory we have to be a bit careful
217 how to not just print() but check whether the object is maybe a pointer to another
218 DataStore object or if it is a TObject with implements Print().
219 Also, call recursively std::pair
220
221 Args:
222 result: object to print
223 depth (int): depth for recursive printing, controls the level of indent
224 weight (float or None): weight to print in addition to object, only used for
225 relations
226 """
227 # are we in recursion? if so indent the output
228 if depth:
229 print(" " * (depth + 1), end="")
230
231 if weight is not None:
232 weight = f" (weight: {weight:.6g})"
233 else:
234 weight = ""
235
236 # is it another RelationsObject? print array name and index
237 if hasattr(result, "getArrayName") and hasattr(result, "getArrayIndex"):
238 if not result:
239 print(f"-> NULL{weight}")
240 else:
241 print(f"-> {result.getArrayName()}#{int(result.getArrayIndex())}{weight}")
242 # special case for TMatrix like types to make them more space efficient
243 elif hasattr(result, "GetNrows") and hasattr(result, "GetNcols"):
244 print(weight, end="")
245 for row in range(result.GetNrows()):
246 print("\n" + " " * (depth + 2), end="")
247 for col in range(result.GetNcols()):
248 print(f"{result(row, col):13.6e} ", end="")
249
250 print()
251 # or is it a TVector3 or TLorentzVector?
252 elif isinstance(result, TVector3):
253 print("(" + ",".join(f"{result[i]:.6g}" for i in range(3)) + ")")
254 elif isinstance(result, XYZVector):
255 print("(" + ",".join(f"{Belle2.B2Vector3D(result)[i]:.6g}" for i in range(3)) + ")")
256 # print("(" + ",".join("%.6g" % x for x in [result.X(), result.Y(), result.Z()]) + ")")
257 elif isinstance(result, TLorentzVector):
258 print("(" + ",".join(f"{result[i]:.6g}" for i in range(4)) + ")")
259 # or, does it look like a std::pair?
260 elif hasattr(result, "first") and hasattr(result, "second"):
261 print(f"pair{weight}")
262 self._printResult(result.first, depth + 1)
263 self._printResult(result.second, depth + 1)
264 # or, could it be a std::vector like container? But ROOT might wrap a std::string so if it has npos assume it's a string
265 elif (hasattr(result, "size") and hasattr(result, "begin")
266 and hasattr(result, "end")) and not hasattr(result, "npos") \
267 and not isinstance(result, Const.DetectorSet):
268 print(f"size({int(result.size())}){weight}")
269 # if it is a RelationVector we also want to print the weights. So
270 # check whether we have weights and pass them to the _printResult
271 weight_getter = getattr(result, "weight", None)
272 weight = None
273 # loop over all elements and print the elements with one level more
274 # indentation
275 for i, e in enumerate(result):
276 if weight_getter is not None:
277 weight = weight_getter(i)
278 self._printResult(e, depth + 1, weight=weight)
279 # print floats with 6 valid digits
280 elif isinstance(result, float):
281 print(f"{result:.6g}{weight}")
282 # print char as int
283 elif isinstance(result, str) and len(result) == 1:
284 print(ord(result), weight, sep="")
285 # ok, in any case we can just print it
286 else:
287 print(result, weight, sep="")
288
289
290# ok, now we just need a small class to call all the printer objects for each
291# event
292class PrintObjectsModule(Module):
294 """Call all DataStorePrinter objects in for each event"""
295
296 def __init__(self, objects_to_print, print_untested=False):
297 """
298 Initialize
299
300 Args:
301 objects_to_print (list): list of object to print
302 """
303
304 self.objects_to_print = objects_to_print
306 self.print_untested = print_untested
307 super().__init__()
308
309 def initialize(self):
310 """Print all untested members if requested"""
311 if not self.print_untested:
312 return
313
314 for printer in self.objects_to_print:
315 printer.print_untested()
316
317 def event(self):
318 """print the contents of the mdst mdst_dataobjects"""
319 try:
320 for printer in self.objects_to_print:
321 printer.print()
322 except Exception as e:
323 B2FATAL("Error in datastore printer: ", e)
324
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72
a (simplified) python wrapper for StoreObjPtr.
Definition: PyStoreObj.h:67
def __init__(self, name, simple, withArgument=None, array=True)
array
if True we print a StoreArray, otherwise a single StoreObjPtr
def add_member(self, name, arguments=None, print_callback=None, display=None)
def _printResult(self, result, depth=0, weight=None)
object_members
list of object members to call and print their results
name
class name of the datastore object
def __init__(self, objects_to_print, print_untested=False)
Abstract base class for different kinds of events.
Definition: ClusterUtils.h:24