diff --git a/addFriend.py b/addFriend.py new file mode 100755 index 0000000000000000000000000000000000000000..ee13395b99e045f19ae59914195ff3482fb03cd6 --- /dev/null +++ b/addFriend.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python + +import argparse + +import ROOT + +parser = argparse.ArgumentParser(description='add a friend tree to a tree in another file') +parser.add_argument("infile", help="input file that contains the friend tree") +parser.add_argument("intree", help="name of the friend tree") +parser.add_argument("outfile", help="output file where the friend tree should be added") +parser.add_argument("outtree", help="name of the tree (in output file) to which the friend should be added") + +args = parser.parse_args() + +outfile = ROOT.TFile.Open(args.outfile, "UPDATE") +infile = ROOT.TFile.Open(args.infile) + +for k in outfile.GetListOfKeys(): + if k.GetName() == args.intree: + raise ValueError("Tree with name {} already exists in outputfile".format(args.intree)) + +outfile.cd() +outtree = outfile.Get(args.outtree) + +if not outtree: + raise KeyError("Tree {} not found in file {}".format(args.outtree, args.outfile)) + + +if outtree.GetListOfFriends(): + for k in outtree.GetListOfFriends(): + if k.GetName() == args.intree: + raise ValueError("Tree with name {} is already friend of {}".format(args.intree, args.outtree)) + +infile.cd() +intree = infile.Get(args.intree) + +if not intree: + raise KeyError("Tree {} not found in file {}".format(args.intree, args.infile)) + +# Add friend and write friend tree and original tree to outfile +outfile.cd() +outtree.AddFriend(intree) +outtree.Write(outtree.GetName()) +outfile.cd() +clonetree = intree.CloneTree(-1, "fast") +clonetree.Write(intree.GetName()) +infile.Close() +outfile.Close()