Make add_friend function importable

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)
-
-friend_name = args.outtree+"_friend_"+args.intree
-
-for k in outfile.GetListOfKeys():
-    if k.GetName() == friend_name:
-        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() == friend_name:
-            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()
-clonetree = intree.CloneTree(-1, "fast")
-clonetree.SetName(friend_name)
-clonetree.SetTitle(friend_name)
-clonetree.Write(friend_name)
-outtree.AddFriend(clonetree)
-outtree.Write(outtree.GetName())
-infile.Close()
-outfile.Close()

add_friend.py

+#!/usr/bin/env python
+
+import ROOT
+
+def add_friend(infile, intree, outfile, outtree):
+
+    root_outfile = ROOT.TFile.Open(outfile, "UPDATE")
+    root_infile = ROOT.TFile.Open(infile)
+
+    friend_name = outtree+"_friend_"+intree
+
+    for k in root_outfile.GetListOfKeys():
+        if k.GetName() == friend_name:
+            raise ValueError("Tree with name {} already exists in outputfile".format(intree))
+
+    root_outfile.cd()
+    root_outtree = root_outfile.Get(outtree)
+
+    if not root_outtree:
+        raise KeyError("Tree {} not found in file {}".format(outtree, outfile))
+
+
+    if root_outtree.GetListOfFriends():
+        for k in root_outtree.GetListOfFriends():
+            if k.GetName() == friend_name:
+                raise ValueError("Tree with name {} is already friend of {}".format(intree, outtree))
+
+    root_infile.cd()
+    root_intree = root_infile.Get(intree)
+
+    if not root_intree:
+        raise KeyError("Tree {} not found in file {}".format(intree, infile))
+
+    # Add friend and write friend tree and original tree to outfile
+    root_outfile.cd()
+    clonetree = root_intree.CloneTree(-1, "fast")
+    clonetree.SetName(friend_name)
+    clonetree.SetTitle(friend_name)
+    clonetree.Write(friend_name)
+    root_outtree.AddFriend(clonetree)
+    root_outtree.Write(root_outtree.GetName())
+    root_infile.Close()
+    root_outfile.Close()
+
+
+if __name__ == "__main__":
+
+    import argparse
+
+    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()
+
+    add_friend(args.infile, args.intree, args.outfile, args.outtree)

toolkit.py

@@ -608,6 +608,7 @@ class ClassificationProject(object):
         f = ROOT.TFile.Open(source_filename)
         tree = f.Get(source_treename)
         entries = tree.GetEntries()
+        logger.info("Write friend tree for {} in {}".format(source_treename, source_filename))
         if os.path.exists(target_filename):
             raise IOError("{} already exists, if you want to recreate it, delete it first".format(target_filename))
         for start in range(0, entries, batch_size):