1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677 |
- #!env/bin/python
- import argparse
- import re
- import os
- import ROOT
- def merge_stl_obj(obj_key, output_file, input1, input_rest, merge_func=None):
- """ Merges STL objects and saves the result into the output file, user
- must supply the merging function.
- """
- obj = input1.Get(obj_key)
- type_name_raw = str(type(obj))
- try:
- type_name = re.findall("<class 'ROOT.([^']+)'>", type_name_raw)[0]
- except IndexError:
- raise ValueError(f"Couldn't extract stl type name from {type_name_raw}")
- if merge_func is not None:
- for input_file in input_rest:
- obj_ = input_file.Get(obj_key)
- merge_func(obj, obj_)
- output_file.WriteObjectAny(obj, type_name, obj_key)
- def merge_obj(obj_key, output_file, input1, input_rest):
- obj = input1.Get(obj_key)
- print('='*80)
- print(f'Merging object {obj_key} of type {type(obj)}')
- if isinstance(obj, ROOT.TH1):
- obj.SetDirectory(output_file) # detach from input file
- for input_file in input_rest:
- obj_ = input_file.Get(obj_key)
- obj.Add(obj_)
- obj.Write()
- else:
- print(f"I don't know how to merge object of type{type(obj)}, but "
- "you can add a case in merge_obj to handle it!")
- def merge_files(input_filenames, output_filename, preserve=False):
- print(f"Merging files {', '.join(input_filenames)} into {output_filename}")
- input1, *input_rest = [ROOT.TFile.Open(input_file, "READ") for input_file in input_filenames]
- output_file = ROOT.TFile.Open(output_filename, "RECREATE")
- output_file.cd()
- obj_keys = [k.GetName() for k in input1.GetListOfKeys()]
- for obj_key in obj_keys:
- if obj_key in {"_value_lookup", "_function_impl_lookup"}:
- merge_stl_obj(obj_key, output_file, input1, [])
- else:
- merge_obj(obj_key, output_file, input1, input_rest)
- for file_ in [input1, *input_rest, output_file]:
- file_.Close()
- print(f"Merge finished! Results have been saved into {output_filename}")
- if preserve:
- print("Preseve specified, leaving input files intact")
- else:
- print("Removing input files...", end='', flush=True)
- for filename in input_filenames:
- os.remove(filename)
- print("Done!")
- if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- add = parser.add_argument
- add('output_file')
- add('input_files', nargs='+')
- add('--preserve', '-p', action='store_true')
- args = parser.parse_args()
- merge_files(args.input_files, args.output_file, args.preserve)
|