merge_files.py

#!/usr/bin/env python3
 

 
import os
import sys
import re
import subprocess
import itertools
from shutil import copyfile
import ROOT
from ROOT.Belle2 import FileMetaData
# we don't really need basf2 but it fixes the print buffering problem
import basf2
from b2test_utils import clean_working_directory, skip_test_if_light
 
 
def create_testfile(name, exp=0, run=0, events=100, branchNames=None, **argk):
    """Create a test file from a steering string"""
    if branchNames is None:
        branchNames = []
    global testfile_steering
    env = dict(os.environ)
    env.update(argk)
 
    steering_file = f"steering-{name}.py"
    with open(steering_file, "w") as f:
        f.write(testfile_steering)
 
    subprocess.call(
        [
            "basf2",
            "-o",
            name,
            "--experiment",
            str(exp),
            "--run",
            str(run),
            "-n",
            str(events),
            steering_file,
        ]
        + branchNames,
        env=env,
    )
 
 
def create_testfile_direct(
    name,
    metadata=None,
    release="test_release",
    user="test_user",
    seed=None,
    site="test_site",
    global_tag="test_globaltag",
    steering="test_steering",
):
    """similar to create_testfile but does it manually without running basf2 for
    full control over the FileMetaData"""
    if metadata is None:
        metadata = FileMetaData()
        if seed is None:
            seed = name + "-seed"
 
    if seed is not None:
        metadata.setRandomSeed(seed)
    metadata.setLfn(name)
    metadata.setCreationData(
        "the most auspicious of days for testing", site, user, release
    )
    metadata.setDatabaseGlobalTag(global_tag)
    metadata.setSteering(steering)
    f = ROOT.TFile(name, "RECREATE")
    t = ROOT.TTree("persistent", "persistent")
    t.Branch("FileMetaData", metadata)
    t.Fill()
    t.Write()
    t = ROOT.TTree("tree", "tree")
    event_meta = ROOT.Belle2.EventMetaData()
    t.Branch("EventMetaData", event_meta)
    t.Fill()
    t.Write()
    f.Close()
 
 
def create_testfile_ntuple(input, output, treeNames=["tree", "anotherTree"], **argk):
    """Create a test ntuple file from a steering string"""
    global testfile_ntuple_steering
    env = dict(os.environ)
    env.update(argk)
 
    steering_file = "steering-ntuple.py"
    with open(steering_file, "w") as f:
        f.write(testfile_ntuple_steering)
 
    subprocess.call(
        ["basf2", "-i", input, "-o", output, steering_file] + treeNames, env=env
    )
 
    # update release in metadata to avoid 'modified-xxx' warnings
    metadata = basf2.get_file_metadata(output)
    metadata.setCreationData(
        metadata.getDate(), metadata.getSite(), metadata.getUser(), "test-release"
    )
    f = ROOT.TFile(output, "UPDATE")
    t = ROOT.TTree("persistent", "persistent")
    t.Branch("FileMetaData", metadata)
    t.Fill()
    t.Write()
    f.Close()
 
 
def merge_files(*args, output="output.root", filter_modified=False):
    """run the merging tool on all passed files
 
    Parameters:
        output: name of the output file
        filter_modified: if True omit warnings that the release is modified and
            consistency cannot be checked
    """
    process = subprocess.run(
        ["b2file-merge", "-q", output] + list(args), stdout=subprocess.PIPE
    )
    # do we want to filter the modified release warning?
    if filter_modified:
        # if so replace them using regular expression
        process.stdout = re.sub(
            rb"^\[WARNING\] File \"(.*?)\" created with modified software ([a-zA-Z0-9\-+]*?): "
            rb"cannot verify that files are compatible\n",
            b"",
            process.stdout,
            flags=re.MULTILINE,
        )
 
    # in any case print output
    sys.stdout.buffer.write(process.stdout)
    sys.stdout.buffer.flush()
    # and return exitcode
    return process.returncode
 
 

testfile_steering = """
import os
import sys
import basf2
basf2.set_log_level(basf2.LogLevel.ERROR)
if "BELLE2_GLOBALTAG" in os.environ:
    basf2.conditions.override_globaltags([os.environ["BELLE2_GLOBALTAG"]])
if "BELLE2_SEED" in os.environ:
    basf2.set_random_seed(os.environ["BELLE2_SEED"])
main = basf2.create_path()
main.add_module("EventInfoSetter")
main.add_module("ParticleGun")
main.add_module("RootOutput", branchNames=sys.argv[1:])
basf2.process(main)
"""
 
 

testfile_ntuple_steering = """
import sys
import basf2
basf2.set_log_level(basf2.LogLevel.ERROR)
main = basf2.create_path()
main.add_module('RootInput')
main.add_module('VariablesToNtuple',
                treeName=sys.argv[1]
                )
main.add_module('VariablesToNtuple',
                treeName=sys.argv[2]
                )
basf2.process(main)
"""
 
 
def check_01_existing():
    """Check that merging a non existing file fails"""
    create_testfile_direct("test2.root")
    return merge_files("/test1.root") != 0 and merge_files("test2.root") == 0
 
 
def check_02_nonroot():
    """Check that merging fails on non-root input files"""
    with open("test1.root", "w") as f:
        f.write("This is not a ROOT file")
    return merge_files("test1.root") != 0
 
 
def check_03_overwrite():
    """Check that overwriting fails if -f is missing"""
    create_testfile_direct("test1.root")
    with open("output.root", "w") as f:
        f.write("stuff")
    return merge_files("test1.root") != 0 and merge_files("-f", "test1.root") == 0
 
 
def check_04_access():
    """Check that it fails if we cannot create output file"""
    create_testfile_direct("test1.root")
    return merge_files("test1.root", output="/nosuchdir/foo") != 0
 
 
def check_05_release():
    """Check that it fails if the releases are different"""
    create_testfile_direct("test1.root")
    create_testfile_direct("test2.root", release="other_release")
    return merge_files("test1.root", "test2.root") != 0
 
 
def check_06_empty_release():
    """Check that merging fails with empty release values"""
    create_testfile_direct("test1.root")
    create_testfile_direct("test2.root", release="")
    return merge_files("test1.root", "test2.root") != 0
 
 
def check_07_modified_release():
    """Check that merging modified release gives warning about that but merging should work"""
    create_testfile_direct("test1.root", release="test_release")
    create_testfile_direct("test2.root", release="test_release-modified")
    return merge_files("test1.root", "test2.root") == 0
 
 
def check_08_duplicate_seed():
    """Check that we get a warning for identical seeds but merging should work"""
    create_testfile_direct("test1.root", seed="seed1")
    create_testfile_direct("test2.root", seed="seed1")
    return merge_files("test1.root", "test2.root") == 0
 
 
def check_09_different_steering():
    """Check that merging fails if the steering file is different"""
    create_testfile_direct("test1.root")
    create_testfile_direct("test2.root", steering="my other steering")
    return merge_files("test1.root", "test2.root") != 0
 
 
def check_10_different_globaltag():
    """Check that merging fails if the global tag is different"""
    create_testfile_direct("test1.root")
    create_testfile_direct("test2.root", global_tag="other_globaltag")
    return merge_files("test1.root", "test2.root") != 0
 
 
def check_11_branches():
    """Check that merging fails if the branches in the event tree are different"""
    create_testfile("test1.root")
    create_testfile("test2.root", branchNames=["EventMetaData"])
    return merge_files("test1.root", "test2.root", filter_modified=True) != 0
 
 
def check_12_hadded():
    """Check that merging fails if the file has more then one entry in the persistent tree"""
    create_testfile_direct("test1.root")
    subprocess.call(["hadd", "test11.root", "test1.root", "test1.root"])
    return merge_files("test11.root") != 0
 
 
def check_13_nopersistent():
    """Check that merging fails without persistent tree"""
    f = ROOT.TFile("test1.root", "RECREATE")
    t = ROOT.TTree("tree", "tree")
    t.Write()
    f.Close()
    return merge_files("test1.root") != 0
 
 
def check_14_noeventtree():
    """Check that merging fails without event tree"""
    f = ROOT.TFile("test1.root", "RECREATE")
    t = ROOT.TTree("persistent", "persistent")
    meta = FileMetaData()
    t.Branch("FileMetaData", meta)
    t.Fill()
    t.Write()
    f.Close()
    return merge_files("test1.root") != 0
 
 
def check_15_noeventbranches():
    """Check that merging fails without event tree"""
    f = ROOT.TFile("test1.root", "RECREATE")
    t = ROOT.TTree("persistent", "persistent")
    meta = FileMetaData()
    meta.setCreationData("date", "site", "user", "release")
    t.Branch("FileMetaData", meta)
    t.Fill()
    t.Write()
    t = ROOT.TTree("tree", "tree")
    t.Write()
    f.Close()
    return merge_files("test1.root") != 0
 
 
def check_16_nonmergeable():
    """Check that merging fails if there are multiple mergeable persistent trees"""
    f = ROOT.TFile("test1.root", "RECREATE")
    t = ROOT.TTree("persistent", "persistent")
    meta = FileMetaData()
    meta.setCreationData("date", "site", "user", "release")
    t.Branch("FileMetaData", meta)
    t.Branch("AnotherMetaData", meta)
    t.Fill()
    t.Write()
    t = ROOT.TTree("tree", "tree")
    t.Branch("EventMetaData", meta)
    t.Fill()
    t.Write()
    f.Close()
    return merge_files("test1.root") != 0
 
 
def check_17_checkparentLFN():
    """Check that parent LFN get merged correctly"""
    parents = [("a", "b", "c"), ("a", "c", "d")]
    m1 = FileMetaData()
    m2 = FileMetaData()
    lfn1 = ROOT.std.vector("std::string")()
    lfn2 = ROOT.std.vector("std::string")()
    for e in parents[0]:
        lfn1.push_back(e)
    for e in parents[1]:
        lfn2.push_back(e)
    m1.setParents(lfn1)
    m2.setParents(lfn2)
    m1.setRandomSeed("1")
    m2.setRandomSeed("2")
    create_testfile_direct("test1.root", m1)
    create_testfile_direct("test2.root", m2)
    merge_files("test1.root", "test2.root")
    meta = basf2.get_file_metadata("output.root")
    should_be = [e for e in sorted(set(parents[0] + parents[1]))]
    is_actual = [meta.getParent(i) for i in range(meta.getNParents())]
    return should_be == is_actual
 
 
def check_18_checkEventNr():
    """Check that event and mc numbers are summed correctly"""
    evtNr = [10, 1243, 232, 1272, 25]
    evtNrFullEvents = [i - 1 for i in evtNr]
    mcNr = [120, 821, 23, 923, 1]
    files = []
    for i, (e, f, m) in enumerate(zip(evtNr, evtNrFullEvents, mcNr)):
        meta = FileMetaData()
        meta.setNEvents(e)
        meta.setNFullEvents(f)
        meta.setMcEvents(m)
        meta.setRandomSeed(str(i))
        files.append(f"test{i}.root")
        create_testfile_direct(files[-1], meta)
    merge_files(*files)
    meta = basf2.get_file_metadata("output.root")
    return (
        sum(evtNr) == meta.getNEvents()
        and sum(evtNrFullEvents) == meta.getNFullEvents()
        and sum(mcNr) == meta.getMcEvents()
    )
 
 
def check_19_lowhigh():
    """Check that the low/high event numbers are merged correctly"""
    lowhigh = [
        (-1, -1, 0),
        (0, 0, 0),
        (0, 0, 1),
        (0, 1, 0),
        (1, 0, 0),
        (1, 1, 1),
    ]
    files = []
    for i, e in enumerate(lowhigh):
        meta = FileMetaData()
        meta.setNEvents(0 if e == (-1, -1, 0) else 1)
        meta.setNFullEvents(0 if e == (-1, -1, 0) else 1)
        meta.setRandomSeed(str(i))
        meta.setLow(e[0], e[1], e[2])
        meta.setHigh(e[0], e[1], e[2])
        files.append(f"test{i}.root")
        create_testfile_direct(files[-1], meta)
 
    # test all possible combinations taking 2 elements from the list plus the
    # full list in one go
    indices = range(len(files))
    tests = list(itertools.permutations(indices, 2)) + [indices]
    for indices in tests:
        low = min(lowhigh[i] for i in indices if lowhigh[i] != (-1, -1, 0))
        high = max(lowhigh[i] for i in indices if lowhigh[i] != (-1, -1, 0))
        if merge_files("-f", "--no-catalog", *(files[i] for i in indices)) != 0:
            return False
        meta = basf2.get_file_metadata("output.root")
        if (
            meta.getExperimentLow() != low[0]
            or meta.getRunLow() != low[1]
            or meta.getEventLow() != low[2]
        ):
            print("low event should be", low)
            meta.Print()
            return False
        if (
            meta.getExperimentHigh() != high[0]
            or meta.getRunHigh() != high[1]
            or meta.getEventHigh() != high[2]
        ):
            print("high event should be", high)
            meta.Print()
            return False
    return True
 
 
def check_20_test_file():
    """Check that a merged file passes the b2file-check program"""
    create_testfile("test1.root", events=1111)
    create_testfile("test2.root", events=123)
    merge_files("test1.root", "test2.root", filter_modified=True)
    return (
        subprocess.call(
            [
                "b2file-check",
                "-n",
                "1234",
                "--mcevents",
                "1234",
                "output.root",
                "EventMetaData",
                "MCParticles",
            ]
        )
        == 0
    )
 
 
def check_21_eventmetadata():
    """Check that merged files has all the correct even infos"""
    create_testfile(
        "test1.root", run=0, events=100, BELLE2_SEED="test1", BELLE2_USER="user1"
    )
    create_testfile(
        "test2.root", run=1, events=100, BELLE2_SEED="test2", BELLE2_USER="user2"
    )
    merge_files("test1.root", "test2.root", "test1.root", filter_modified=True)
    out = ROOT.TFile("output.root")
    events = out.Get("tree")
    entries = events.GetEntriesFast()
    if entries != 300:
        return False
    # we expect to see the events from run 0 twice and the ones from run 1 once.
    # So create a dictionary which contains the expected counts
    eventcount = {(0, 0, i + 1): 2 for i in range(100)}
    eventcount.update({(0, 1, i + 1): 1 for i in range(100)})
    for i in range(entries):
        events.GetEntry(i)
        e = events.EventMetaData
        eventcount[(e.getExperiment(), e.getRun(), e.getEvent())] -= 1
    return max(eventcount.values()) == 0 and min(eventcount.values()) == 0
 
 
def check_22_real_mc():
    """Check that merging fails if real and MC data are mixed"""
    create_testfile_direct("test1.root")
    copyfile(basf2.find_file("framework/tests/fake_real.root"), "test2.root")
    return merge_files("test1.root", "test2.root") != 0
 
 
def check_23_legacy_ip():
    """Check that we can merge if the Legacy_IP_Information is inconsistent"""
    create_testfile_direct("test1.root", global_tag="test_globaltag")
    create_testfile_direct(
        "test2.root", global_tag="test_globaltag,Legacy_IP_Information"
    )
    if merge_files("test1.root", "test2.root") != 0:
        return False
    meta = basf2.get_file_metadata("output.root")
    return meta.getDatabaseGlobalTag() == "test_globaltag"
 
 
def check_24_legacy_ip_middle():
    """Check that we can merge if the Legacy_IP_Information is inconsistent"""
    create_testfile_direct("test1.root", global_tag="test_globaltag,other")
    create_testfile_direct(
        "test2.root", global_tag="test_globaltag,Legacy_IP_Information,other"
    )
    if merge_files("test1.root", "test2.root") != 0:
        return False
    meta = basf2.get_file_metadata("output.root")
    return meta.getDatabaseGlobalTag() == "test_globaltag,other"
 
 
def check_25_legacy_ip_only():
    """Check that we can merge if the Legacy_IP_Information is inconsistent"""
    create_testfile_direct("test1.root", global_tag="")
    create_testfile_direct("test2.root", global_tag="Legacy_IP_Information")
    if merge_files("test1.root", "test2.root") != 0:
        return False
    meta = basf2.get_file_metadata("output.root")
    return meta.getDatabaseGlobalTag() == ""
 
 
def check_26_ntuple_merge():
    """Check that we can merge two ntuple output files"""
    create_testfile("test1.root", exp=1, run=2, events=111)
    create_testfile("test2.root", exp=1, run=2, events=123)
    create_testfile_ntuple(input="test1.root", output="ntuple1.root")
    create_testfile_ntuple(input="test2.root", output="ntuple2.root")
    return merge_files("ntuple1.root", "ntuple2.root") == 0
 
 
def check_27_ntuple_trees():
    """Check that ntuple merge fails if the tree names are different"""
    create_testfile("test1.root")
    create_testfile("test2.root")
    create_testfile_ntuple(input="test1.root", output="ntuple1.root")
    create_testfile_ntuple(
        input="test2.root", output="ntuple2.root", treeNames=["differentTree", "tree"]
    )
    return merge_files("ntuple1.root", "ntuple2.root") != 0
 
 
def check_28_streaming():
    """Check if we can merge streamed input files"""
    # Here we use as input a mdst file from GitHub
    input_file = 'https://github.com/belle2/basf2/raw/refs/heads/main/mdst/tests/mdst-v09-00-00.root'
    return merge_files(input_file) == 0
 
 
def check_29_parent_release():
    """Check that merging files does not modify the release version in the metadata."""
    create_testfile_direct("test1.root", release="abcd")
    create_testfile_direct("test2.root", release="abcd")
    merge_files("test1.root", "test2.root", output="output.root")
    meta = basf2.get_file_metadata("output.root")
    return meta.getRelease() == "abcd"
 
 
def check_XX_filemetaversion():
    """Check that the Version of the FileMetaData hasn't changed.
    If this check fails please check that the changes to FileMetaData don't
    affect b2file-merge and adapt the correct version number here."""
    return FileMetaData.Class().GetClassVersion() == 11
 
 
if __name__ == "__main__":
    skip_test_if_light()  # light builds don't have particle gun
    failures = 0
    existing = [e for e in sorted(globals().items()) if e[0].startswith("check_")]
    for name, fcn in existing:
        print(f"running {name}: {fcn.__doc__}")
        with clean_working_directory():
            if not fcn():
                print(f"{name} failed")
                failures += 1
            else:
                print(f"{name} passed")
 
    sys.exit(failures)