gen_flatrecord.py
Go to the documentation of this file.
1 #!/usr/bin/env python
2 
3 # This script is intended to be compatible with both python2 and python3
4 from __future__ import print_function
5 
6 from pygccxml import *
7 
8 import os
9 import sys
10 
11 # Types that ROOT will understand (with a little nudging in some cases)
12 fundamental_types = ['int', 'float', 'double', 'bool', 'unsigned int',
13  'short', 'short int', 'short unsigned int',
14  'long', 'long unsigned int',
15  'long long int', 'unsigned char',
16  'size_t',
17  'std::string']
18 
19 enums = []
20 
21 def is_fundamental(type):
22  while type[:5] == 'caf::': type = type[5:]
23 
24  return type in fundamental_types or type in enums
25 
26 
27 def translate_type(type):
28  while type[:5] == 'caf::': type = type[5:]
29 
30  # enums are small, right?
31  if type in enums: return 'short unsigned int'
32  # TTree can't handle long long at all? This will lose information though...
33  if type == 'long long int': return 'long'
34 
35  return type
36 
37 
38 def is_vector(type):
39  return type[:12] == 'std::vector<'
40 
41 
42 def vector_contents(type):
43  assert is_vector(type)
44  return type[12:type.find(',')]
45 
46 
47 def is_array(type):
48  return '[' in type
49 
50 
51 def array_type(type):
52  assert is_array(type)
53  # There always seems to be a space before the dimensions
54  return type[:type.find('[')-1]
55 
56 
57 def array_size(type):
58  assert is_array(type)
59  return int(type[type.find('[')+1:-1])
60 
61 
63  while type[:5] == 'caf::': type = type[5:]
64 
65  assert not is_fundamental(type)
66  assert not is_vector(type)
67 
68  if type == 'StandardRecord': return 'FlatRecord'
69 
70  return 'Flat'+type[2:]
71 
72 
73 def base_class(klass):
74  assert len(klass.bases) < 2, 'Support for multiple base classes not implemented'
75  if len(klass.bases) == 1:
76  return klass.bases[0].related_class
77  return None
78 
79 
80 # Recurse to find member variables including all base classes
82  vars = list(klass.variables())
83  base = base_class(klass)
84  if base: vars += variables_inc_bases(base)
85  return vars
86 
87 
88 def branch_cmd(name, treeName = 'tr'):
89  return 'branch('+treeName+', prefix+"'+name+'", &'+name+', policy);'
90 
91 def branch_cmd_no_policy(name, treeName = 'tr'):
92  return 'branch('+treeName+', prefix+"'+name+'", &'+name+', 0);'
93 
94 
95 
96 if len(sys.argv) < 1 or len(sys.argv) > 3:
97  print('Usage: gen_flatrecord.py [/path/to/header/outputs/] [/path/to/cxx/outputs/]')
98  sys.exit(1)
99 
100 headerDir = os.getcwd()
101 if len(sys.argv) >= 2: headerDir = sys.argv[1]
102 cxxDir = os.getcwd()
103 if len(sys.argv) >= 3: cxxDir = sys.argv[2]
104 
105 # Locate the castxml executable
106 generator_path, generator_name = utils.find_xml_generator()
107 
108 # Figure out where our source files are
109 priv = os.environ['SRT_PRIVATE_CONTEXT'] if 'SRT_PRIVATE_CONTEXT' in os.environ else None
110 pub = os.environ['SRT_PUBLIC_CONTEXT'] if 'SRT_PUBLIC_CONTEXT' in os.environ else None
111 
112 # For cmake build
113 if not pub: pub = os.environ['NOVASOFT_DIR']
114 
115 context = priv if os.path.exists(priv+'/StandardRecord/') else pub
116 
117 path = []
118 if priv: path += [priv]
119 path += [pub]
120 path += [os.environ['ROOT_INC']]
121 
122 config = parser.xml_generator_configuration_t(
123  xml_generator_path=generator_path,
124  xml_generator=generator_name,
125  include_paths=path,
126  cflags='-std=c++1z -Wno-unknown-warning-option'#,
127 # start_with_declarations='caf::StandardRecord'
128  )
129 
130 decls = parser.parse([context+'/StandardRecord/StandardRecord.h'],
131  config)
132 
133 global_namespace = declarations.get_global_namespace(decls)
134 ns = global_namespace.namespace('caf')
135 
136 enums += [e.name for e in ns.enumerations()]
137 
138 # Keep track of which classes we've written out so far, for purposes of
139 # dependency tracking.
140 emitted = []
141 
142 
143 disclaimer = '''// This file was auto-generated by gen_flatrecord.py.
144 // DO NOT EDIT IT DIRECTLY.
145 // For documentation of the fields see the regular StandardRecord.h'''
146 
147 # We write individual headers so that the preprocessor will handle our
148 # dependency order, but put all the code in one .cxx becase it's much faster to
149 # compile that way.
150 sys.stdout = open(cxxDir+'/FlatRecord.cxx', 'wa')
151 print(disclaimer)
152 print()
153 print('#include "StandardRecord/StandardRecord.h"')
154 print()
155 print('#include "FlatMaker/TreeUtils.h"')
156 print()
157 print('#include "TTree.h"')
158 
159 for klass in ns.classes():
160  pt = type_to_flat_type(klass.name)
161 
162  sys.stdout = open(headerDir+'/'+pt+'.h', 'w')
163 
164  class TypeName:
165  def __init__(self, t, n, length = -1):
166  self.type = t
167  self.name = n
168  self.length = length
169 
170  def is_array(self): return self.length > 0
171 
172  basic_mems = []
173  class_mems = []
174  vec_mems = []
175  for v in variables_inc_bases(klass):
176  type = str(v.decl_type)
177 
178  if is_fundamental(type):
179  basic_mems += [TypeName(translate_type(type), v.name)]
180  elif is_array(type):
181  innert = array_type(type)
182  size = array_size(type)
183  if is_fundamental(innert):
184  vec_mems += [TypeName(innert, v.name, size)]
185  else:
186  vec_mems += [TypeName(type_to_flat_type(innert), v.name, size)]
187  elif is_vector(type):
188  innert = vector_contents(type)
189  if is_fundamental(innert):
190  vec_mems += [TypeName(innert, v.name)]
191  else:
192  vec_mems += [TypeName(type_to_flat_type(innert), v.name)]
193  else:
194  class_mems += [TypeName(type_to_flat_type(type), v.name)]
195 
196  print(disclaimer)
197  print()
198  print('#pragma once')
199  print()
200  print('#include "StandardRecord/SREnums.h"')
201  print()
202  incs = set()
203  for tn in class_mems: incs.add(tn.type)
204  for tn in vec_mems:
205  if not is_fundamental(tn.type):
206  incs.add(tn.type)
207  for i in incs: print('#include "FlatMaker/'+i+'.h"')
208  print()
209  print('#include <string>')
210  print()
211  print('class TTree;')
212  print('namespace caf{class '+klass.name+';}')
213  print()
214  print('namespace flat')
215  print('{')
216  print('class IBranchPolicy;')
217  print()
218  print('/// Flat encoding of \\ref', klass.name)
219  print('class', pt)
220  print('{')
221  print('public:')
222  print(' '+pt+'(const std::string& prefix, TTree* tr, const IBranchPolicy* policy);')
223  print(' ~'+pt+'();')
224  print()
225  print(' void Fill(const caf::'+klass.name+'& sr);')
226  print()
227  print('protected:')
228 
229  # Members
230  for tn in basic_mems:
231  print(' '+tn.type+' '+tn.name+';')
232 
233  if len(basic_mems) > 0 and len(class_mems) > 0: print()
234 
235  for tn in class_mems:
236  print(' '+tn.type+' '+tn.name+';')
237 
238  for tn in vec_mems:
239  print()
240  print(' TTree* '+tn.name+'_tree;')
241  print(' '+tn.type+' '+tn.name+';')
242  print(' long '+tn.name+'_idx;')
243  print(' int '+tn.name+'_length;')
244 
245  print('};')
246  print('} // end namespace')
247 
248  # Switch to appending into the main cxx file
249  sys.stdout = open(cxxDir+'/FlatRecord.cxx', 'a')
250 
251  print()
252  print('#include "FlatMaker/'+pt+'.h"')
253  print()
254  print('flat::'+pt+'::'+pt+'(const std::string& prefix, TTree* tr, const IBranchPolicy* policy)')
255 
256  # Initializer list
257  inits = [tn.name+'(prefix+"'+tn.name+'.", tr, policy)' for tn in class_mems]
258 
259  for tn in vec_mems:
260  inits += [tn.name+'_tree(make_tree(prefix+"'+tn.name+'", "'+tn.name+'", tr))']
261 
262  if is_fundamental(tn.type):
263  inits += [tn.name+'(0)']
264  else:
265  inits += [tn.name+'((prefix+"'+tn.name+'."), '+tn.name+'_tree, policy)']
266 
267  if tn.is_array():
268  inits += [tn.name+'_idx(0), '+tn.name+'_length(sizeof(caf::'+klass.name+'::'+tn.name+')/sizeof(*caf::'+klass.name+'::'+tn.name+'))']
269  else:
270  inits += [tn.name+'_idx(0), '+tn.name+'_length(0)']
271 
272  if len(inits) > 0: print(' : '+',\n '.join(inits))
273 
274  print('{')
275 
276  # Constructor body
277  for tn in basic_mems:
278  print(' '+branch_cmd(tn.name))
279 
280  for tn in vec_mems:
281  if is_fundamental(tn.type):
282  print(' '+branch_cmd(tn.name, tn.name+'_tree'))
283  print(' if('+tn.name+'_tree->GetNbranches() > 0){')
284  print(' '+branch_cmd_no_policy(tn.name+'_idx'))
285  print(' '+branch_cmd_no_policy(tn.name+'_length'))
286  print(' }')
287 
288  print('}')
289 
290  print()
291 
292  print('flat::'+pt+'::~'+pt+'()')
293  print('{')
294  for tn in vec_mems:
295  trName = tn.name+'_tree'
296  print(' if('+trName+'->GetNbranches() > 0) '+trName+'->Write();')
297  print(' delete '+trName+';')
298  print('}')
299 
300  print()
301 
302  print('void flat::'+pt+'::Fill(const caf::'+klass.name+'& sr)')
303  print('{')
304  for tn in basic_mems:
305  print(' '+tn.name+' = sr.'+tn.name+';')
306 
307  if len(basic_mems) > 0 and len(class_mems) > 0: print()
308 
309  for tn in class_mems:
310  print(' '+tn.name+'.Fill(sr.'+tn.name+');')
311 
312  for tn in vec_mems:
313  print()
314  print(' '+tn.name+'_idx += '+tn.name+'_length; // increment taken by previous record')
315  if not tn.is_array(): # have to re-fill each time for vectors
316  print(' '+tn.name+'_length = sr.'+tn.name+'.size();')
317  print(' for(const auto& x: sr.'+tn.name+'){')
318  if is_fundamental(tn.type):
319  print(' '+tn.name+' = x;')
320  else:
321  print(' '+tn.name+'.Fill(x);')
322  print(' '+tn.name+'_tree->Fill();');
323  print(' }')
324  print('}')
325 
326 sys.stderr.write('Wrote FlatRecord.cxx and associated headers\n')
def base_class(klass)
def variables_inc_bases(klass)
def array_type(type)
def branch_cmd_no_policy(name, treeName='tr')
def array_size(type)
def branch_cmd(name, treeName='tr')
bool print
def is_fundamental(type)
def is_array(type)
procfile open("FD_BRL_v0.txt")
def translate_type(type)
def __init__(self, t, n, length=-1)
def is_vector(type)
def vector_contents(type)
def type_to_flat_type(type)