Adding script to add friend trees to files

addFriend.py

+#!/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()