merge.py 2.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677
  1. #!env/bin/python
  2. import argparse
  3. import re
  4. import os
  5. import ROOT
  6. def merge_stl_obj(obj_key, output_file, input1, input_rest, merge_func=None):
  7. """ Merges STL objects and saves the result into the output file, user
  8. must supply the merging function.
  9. """
  10. obj = input1.Get(obj_key)
  11. type_name_raw = str(type(obj))
  12. try:
  13. type_name = re.findall("<class 'ROOT.([^']+)'>", type_name_raw)[0]
  14. except IndexError:
  15. raise ValueError(f"Couldn't extract stl type name from {type_name_raw}")
  16. if merge_func is not None:
  17. for input_file in input_rest:
  18. obj_ = input_file.Get(obj_key)
  19. merge_func(obj, obj_)
  20. output_file.WriteObjectAny(obj, type_name, obj_key)
  21. def merge_obj(obj_key, output_file, input1, input_rest):
  22. obj = input1.Get(obj_key)
  23. print('='*80)
  24. print(f'Merging object {obj_key} of type {type(obj)}')
  25. if isinstance(obj, ROOT.TH1):
  26. obj.SetDirectory(output_file) # detach from input file
  27. for input_file in input_rest:
  28. obj_ = input_file.Get(obj_key)
  29. obj.Add(obj_)
  30. obj.Write()
  31. else:
  32. print(f"I don't know how to merge object of type{type(obj)}, but "
  33. "you can add a case in merge_obj to handle it!")
  34. def merge_files(input_filenames, output_filename, preserve=False):
  35. print(f"Merging files {', '.join(input_filenames)} into {output_filename}")
  36. input1, *input_rest = [ROOT.TFile.Open(input_file, "READ") for input_file in input_filenames]
  37. output_file = ROOT.TFile.Open(output_filename, "RECREATE")
  38. output_file.cd()
  39. obj_keys = [k.GetName() for k in input1.GetListOfKeys()]
  40. for obj_key in obj_keys:
  41. if obj_key in {"_value_lookup", "_function_impl_lookup"}:
  42. merge_stl_obj(obj_key, output_file, input1, [])
  43. else:
  44. merge_obj(obj_key, output_file, input1, input_rest)
  45. for file_ in [input1, *input_rest, output_file]:
  46. file_.Close()
  47. print(f"Merge finished! Results have been saved into {output_filename}")
  48. if preserve:
  49. print("Preseve specified, leaving input files intact")
  50. else:
  51. print("Removing input files...", end='', flush=True)
  52. for filename in input_filenames:
  53. os.remove(filename)
  54. print("Done!")
  55. if __name__ == '__main__':
  56. parser = argparse.ArgumentParser()
  57. add = parser.add_argument
  58. add('output_file')
  59. add('input_files', nargs='+')
  60. add('--preserve', '-p', action='store_true')
  61. args = parser.parse_args()
  62. merge_files(args.input_files, args.output_file, args.preserve)