import ROOT
import sys

# A simple helper function to fill a test tree: this makes the example stand-alone.
def fill_tree(treeName, fileName, b1, b2):
   opts = ROOT.ROOT.RDF.RSnapshotOptions();
   opts.fMode = "UPDATE";

   branchList = ROOT.vector('string')()
   for branchName in [b1, b2]:
     branchList.push_back(branchName)
     
   tdf = ROOT.ROOT.RDataFrame(200)
   tdf.Define(b1, "(double) tdfentry_")\
      .Define(b2, "(int) tdfentry_ * tdfentry_")\
      .Snapshot(treeName, fileName, branchList, opts)

# We prepare some trees split on different files
for f in range(20):
   filename = "sample{}.root".format(f)
   fill_tree("processInfoTree", filename, "luminosity", "cross_section")
   fill_tree("weightsTree",     filename, "weight", "weight1")
   fill_tree("randomTree",      filename, "x", "x2")
   fill_tree("anotherTree",     filename, "y", "y2")