def test_iterate_report_2(): with uproot4.open( skhep_testdata.data_path("uproot-HZZ.root"))["events"] as events: for i, (arrays, report) in enumerate( events.iterate("Muon_Px", step_size=1000, report=True, library="np")): if i == 0: assert report.tree_entry_start == 0 assert report.tree_entry_stop == 1000 assert report.file_path == skhep_testdata.data_path( "uproot-HZZ.root") elif i == 1: assert report.tree_entry_start == 1000 assert report.tree_entry_stop == 2000 assert report.file_path == skhep_testdata.data_path( "uproot-HZZ.root") elif i == 2: assert report.tree_entry_start == 2000 assert report.tree_entry_stop == 2421 assert report.file_path == skhep_testdata.data_path( "uproot-HZZ.root") else: assert False
def test_iofeatures_offsets(): with uproot.open( skhep_testdata.data_path("uproot-small-dy-withoffsets.root") )["tree/Muon_pt"] as withoffsets: muonpt1 = withoffsets.array(library="np", entry_start=10, entry_stop=20) assert [x.tolist() for x in muonpt1] == [ [20.60145378112793], [50.36957550048828, 41.21387481689453, 3.1869382858276367], [51.685970306396484, 35.227813720703125], [], [], [], [], [23.073759078979492], [32.921417236328125, 8.922308921813965, 4.368383407592773], [51.9132194519043, 31.930095672607422], ] with uproot.open(skhep_testdata.data_path( "uproot-small-dy-nooffsets.root"))["tree/Muon_pt"] as nooffsets: muonpt2 = nooffsets.array(library="np", entry_start=10, entry_stop=20) assert [x.tolist() for x in muonpt2] == [ [20.60145378112793], [50.36957550048828, 41.21387481689453, 3.1869382858276367], [51.685970306396484, 35.227813720703125], [], [], [], [], [23.073759078979492], [32.921417236328125, 8.922308921813965, 4.368383407592773], [51.9132194519043, 31.930095672607422], ]
def test_lazy_colon(): uproot4.lazy(skhep_testdata.data_path("uproot-issue63.root") + ":WtLoop_nominal") uproot4.lazy( [ skhep_testdata.data_path("uproot-issue63.root") + ":WtLoop_nominal", skhep_testdata.data_path("uproot-issue63.root") + ":WtLoop_Fake_nominal", ] )
def test_data_path_missing(): path = skhtd.data_path("doesnt-exist.root", raise_missing=False) assert path == os.path.join(data_dir, "doesnt-exist.root") with pytest.raises(RuntimeError): skhtd.data_path("doesnt-exist.root") with pytest.raises(RuntimeError): skhtd.data_path("doesnt-exist.root", raise_missing=True)
def test_recovery(mini): # flat array to recover: filename = skhep_testdata.data_path("uproot-issue21.root") with uproot.open({"file:" + filename: "nllscan/mH"}, minimal_ttree_metadata=mini) as branch: basket = branch.basket(0) assert basket.data.view(">f8").tolist()[:10] == [ 124.0, 124.09089660644531, 124.18180084228516, 124.27269744873047, 124.36360168457031, 124.45449829101562, 124.54550170898438, 124.63639831542969, 124.72730255126953, 124.81819915771484, ] assert basket.byte_offsets is None assert branch.entry_offsets == [0, branch.num_entries] # jagged arrays to recover: # uproot-issue327.root DstTree: fTracks.fCharge # uproot-issue232.root fTreeV0: V0s.fV0pt MCparticles.nbodies # uproot-issue187.root fTreeV0: V0s.fV0pt MCparticles.nbodies # uproot-from-geant4.root Details: numgood, TrackedRays: Event phi filename = skhep_testdata.data_path("uproot-issue327.root") with uproot.open({"file:" + filename: "DstTree/fTracks.fCharge"}, minimal_ttree_metadata=mini) as branch: basket = branch.basket(0) assert basket.data.view("i1")[:10].tolist() == [ 1, -1, 1, 1, -1, -1, 1, -1, -1, -1, ] assert basket.byte_offsets[:10].tolist() == [ 0, 2, 37, 56, 60, 81, 82, 112, 112, 112, ] assert branch.entry_offsets == [0, branch.num_entries]
def test_fallback_reading(): with uproot4.open( skhep_testdata.data_path("uproot-small-evnt-tree-fullsplit.root") ) as f: f["tree:evt/P3/P3.Py"] assert f.file._streamers is None with uproot4.open(skhep_testdata.data_path("uproot-demo-double32.root")) as f: f["T/fD64"] assert f.file._streamers is not None
def test_multiprocessing(): pool = multiprocessing.Pool(1) out = pool.map( readone, [ skhep_testdata.data_path("uproot-Zmumu.root"), skhep_testdata.data_path("uproot-Zmumu-zlib.root"), ], ) list(out)
def test_function_iterate_pandas_2(): pandas = pytest.importorskip("pandas") files = [ skhep_testdata.data_path("uproot-HZZ.root") + ":events", skhep_testdata.data_path("uproot-HZZ-uncompressed.root") + ":events", skhep_testdata.data_path("uproot-HZZ-zlib.root") + ":events", skhep_testdata.data_path("uproot-HZZ-lz4.root") + ":events", ] expect = 0 for arrays, report in uproot.iterate(files, "Muon_Px", report=True, library="pd"): assert arrays["Muon_Px"].index.values[0] == (expect, 0) expect += report.tree.num_entries
def test_lazy(): with pytest.raises(ValueError): uproot4.lazy(skhep_testdata.data_path("uproot-issue63.root")) with pytest.raises(ValueError): uproot4.lazy( {skhep_testdata.data_path("uproot-issue63.root"): "blah"}, allow_missing=True, ) uproot4.lazy({skhep_testdata.data_path("uproot-issue63.root"): "WtLoop_nominal"}) uproot4.lazy( { skhep_testdata.data_path("uproot-issue63.root"): "WtLoop_nominal", skhep_testdata.data_path("uproot-issue63.root"): "WtLoop_Fake_nominal", } ) uproot4.lazy([{skhep_testdata.data_path("uproot-issue63.root"): "WtLoop_nominal"}]) uproot4.lazy( {skhep_testdata.data_path("uproot-issue63.root") + "*": "WtLoop_nominal"} ) uproot4.lazy( [{skhep_testdata.data_path("uproot-issue63.root") + "*": "WtLoop_nominal"}] )
def test_open(): assert isinstance( uproot4.open(skhep_testdata.data_path("uproot-issue63.root")), uproot4.reading.ReadOnlyDirectory, ) assert isinstance( uproot4.open( {skhep_testdata.data_path("uproot-issue63.root"): "WtLoop_nominal"} ), uproot4.behaviors.TTree.TTree, ) with pytest.raises(ValueError): uproot4.open([skhep_testdata.data_path("uproot-issue63.root")])
def test(): one = skhep_testdata.data_path("uproot-sample-6.16.00-uncompressed.root") two = skhep_testdata.data_path("uproot-sample-6.18.00-uncompressed.root") bad = one.replace(".root", "-DOES-NOT-EXIST.root") okay = one.replace(".root", "-DOES-NOT-EXIST-*.root") assert len(list(uproot.iterate([one, two], step_size="1 TB", library="np"))) == 2 with pytest.raises(uproot._util._FileNotFoundError): list(uproot.iterate([one, two, bad], library="np")) assert (len( list(uproot.iterate([one, two, okay], step_size="1 TB", library="np"))) == 2)
def test_cache(): with uproot4.open( skhep_testdata.data_path( "uproot-sample-6.20.04-uncompressed.root"), object_cache=100, array_cache="100 MB", ) as f: assert f.cache_key == "db4be408-93ad-11ea-9027-d201a8c0beef:/" assert f[ "sample"].cache_key == "db4be408-93ad-11ea-9027-d201a8c0beef:/sample;1" assert (f["sample/i4"].cache_key == "db4be408-93ad-11ea-9027-d201a8c0beef:/sample;1:i4(16)") i4 = f["sample/i4"] assert list(f.file.array_cache) == [] i4.array(uproot4.interpretation.numerical.AsDtype(">i4"), library="np") assert list(f.file.array_cache) == [ "db4be408-93ad-11ea-9027-d201a8c0beef:/sample;1:i4(16):AsDtype(Bi4(),Li4()):0-30:np" ] with pytest.raises(OSError): i4.array(uproot4.interpretation.numerical.AsDtype(">i4"), entry_start=3, library="np") i4.array(uproot4.interpretation.numerical.AsDtype(">i4"), library="np")
def test_branch_array_4(file_handler): with uproot.open( skhep_testdata.data_path("uproot-sample-6.20.04-uncompressed.root"), file_handler=file_handler, )["sample/i4"] as branch: with pytest.raises(ValueError): branch.array(uproot.interpretation.numerical.AsDtype(">i8"), library="np")
def test_branch_array_2(file_handler): with uproot.open( skhep_testdata.data_path("uproot-sample-6.20.04-uncompressed.root"), file_handler=file_handler, )["sample/i4"] as branch: assert branch.array( uproot.interpretation.numerical.AsDtype(">i4"), entry_start=3, entry_stop=-5, library="np", ).tolist() == [ -12, -11, -10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, ]
def test_ranges_or_baskets_to_arrays(): with uproot.open( skhep_testdata.data_path("uproot-sample-6.20.04-uncompressed.root") )["sample"] as sample: branch = sample["i4"] ranges_or_baskets = _names_entries_to_ranges_or_baskets(sample, ["i4"], 0, 30) branchid_interpretation = { branch.cache_key: uproot.interpretation.numerical.AsDtype(">i4") } entry_start, entry_stop = (0, 30) decompression_executor = uproot.source.futures.TrivialExecutor() interpretation_executor = uproot.source.futures.TrivialExecutor() library = uproot.interpretation.library._libraries["np"] arrays = {} uproot.behaviors.TBranch._ranges_or_baskets_to_arrays( sample, ranges_or_baskets, branchid_interpretation, entry_start, entry_stop, decompression_executor, interpretation_executor, library, arrays, False, ) assert arrays[branch.cache_key].tolist() == [ -15, -14, -13, -12, -11, -10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, ]
def test_map_int32_vector_vector_int16(): with uproot4.open(skhep_testdata.data_path( "uproot-stl_containers.root"))["tree"] as tree: assert [ x.tolist() for x in tree["map_int32_vector_vector_int16"].array(library="np") ] == [ { 1: [[1]] }, { 1: [[1]], 2: [[1], [1, 2]] }, { 1: [[1]], 2: [[1], [1, 2]], 3: [[1], [1, 2], [1, 2, 3]] }, { 1: [[1]], 2: [[1], [1, 2]], 3: [[1], [1, 2], [1, 2, 3]], 4: [[1], [1, 2], [1, 2, 3], [1, 2, 3, 4]], }, { 1: [[1]], 2: [[1], [1, 2]], 3: [[1], [1, 2], [1, 2, 3]], 4: [[1], [1, 2], [1, 2, 3], [1, 2, 3, 4]], 5: [[1], [1, 2], [1, 2, 3], [1, 2, 3, 4], [1, 2, 3, 4, 5]], }, ]
def test_map_string_vector_int16(): with uproot4.open(skhep_testdata.data_path( "uproot-stl_containers.root"))["tree"] as tree: assert [ x.tolist() for x in tree["map_string_vector_int16"].array(library="np") ] == [ { "one": [1] }, { "one": [1], "two": [1, 2] }, { "one": [1], "two": [1, 2], "three": [1, 2, 3] }, { "one": [1], "two": [1, 2], "three": [1, 2, 3], "four": [1, 2, 3, 4] }, { "one": [1], "two": [1, 2], "three": [1, 2, 3], "four": [1, 2, 3, 4], "five": [1, 2, 3, 4, 5], }, ]
def test_map_int32_set_int16(): with uproot4.open(skhep_testdata.data_path( "uproot-stl_containers.root"))["tree"] as tree: assert [ x.tolist() for x in tree["map_int32_set_int16"].array(library="np") ] == [ { 1: set([1]) }, { 1: set([1]), 2: set([1, 2]) }, { 1: set([1]), 2: set([1, 2]), 3: set([1, 2, 3]) }, { 1: set([1]), 2: set([1, 2]), 3: set([1, 2, 3]), 4: set([1, 2, 3, 4]) }, { 1: set([1]), 2: set([1, 2]), 3: set([1, 2, 3]), 4: set([1, 2, 3, 4]), 5: set([1, 2, 3, 4, 5]), }, ]
def test_map_int32_int16(): with uproot4.open(skhep_testdata.data_path( "uproot-stl_containers.root"))["tree"] as tree: assert [ x.tolist() for x in tree["map_int32_int16"].array(library="np") ] == [ { 1: 1 }, { 1: 1, 2: 2 }, { 1: 1, 2: 2, 3: 3 }, { 1: 1, 2: 2, 3: 3, 4: 4 }, { 1: 1, 2: 2, 3: 3, 4: 4, 5: 5 }, ]
def test_file_header(): filename = skhep_testdata.data_path("uproot-Zmumu.root") file = uproot4.reading.ReadOnlyFile(filename) assert repr(file.compression) == "ZLIB(4)" assert not file.is_64bit assert file.fNbytesInfo == 4447 assert file.hex_uuid == "944b77d0-98ab-11e7-a769-0100007fbeef"
def test_vector_set_string(): with uproot4.open(skhep_testdata.data_path( "uproot-stl_containers.root"))["tree"] as tree: assert [ x.tolist() for x in tree["vector_set_string"].array(library="np") ] == [ [set(["one"])], [set(["one"]), set(["one", "two"])], [set(["one"]), set(["one", "two"]), set(["one", "two", "three"])], [ set(["one"]), set(["one", "two"]), set(["one", "two", "three"]), set(["one", "two", "three", "four"]), ], [ set(["one"]), set(["one", "two"]), set(["one", "two", "three"]), set(["one", "two", "three", "four"]), set(["one", "two", "three", "four", "five"]), ], ]
def test_new_name(tmp_path): newfile = os.path.join(tmp_path, "newfile.root") with uproot.open( skhep_testdata.data_path("uproot-histograms.root")) as fin: one = fin["one"] with uproot.recreate(newfile) as fout: fout["whatever"] = one f1 = ROOT.TFile(newfile) h1 = f1.Get("whatever") assert h1.GetBinContent(0) == 0 assert h1.GetBinContent(1) == 68 assert h1.GetBinContent(2) == 285 assert h1.GetBinContent(3) == 755 assert h1.GetBinContent(4) == 1580 assert h1.GetBinContent(5) == 2296 assert h1.GetBinContent(6) == 2286 assert h1.GetBinContent(7) == 1570 assert h1.GetBinContent(8) == 795 assert h1.GetBinContent(9) == 289 assert h1.GetBinContent(10) == 76 assert h1.GetBinContent(11) == 0 f1.Close()
def test_map_int32_vector_string(): with uproot4.open(skhep_testdata.data_path( "uproot-stl_containers.root"))["tree"] as tree: assert [ x.tolist() for x in tree["map_int32_vector_string"].array(library="np") ] == [ { 1: ["one"] }, { 1: ["one"], 2: ["one", "two"] }, { 1: ["one"], 2: ["one", "two"], 3: ["one", "two", "three"] }, { 1: ["one"], 2: ["one", "two"], 3: ["one", "two", "three"], 4: ["one", "two", "three", "four"], }, { 1: ["one"], 2: ["one", "two"], 3: ["one", "two", "three"], 4: ["one", "two", "three", "four"], 5: ["one", "two", "three", "four", "five"], }, ]
def test(): with uproot.open(skhep_testdata.data_path("uproot-histograms.root")) as f: streamers = f.file.streamers assert len(streamers) == 14 assert drop_fbits(from_ROOT["TH1F"]) == drop_fbits( streamers["TH1F"][2].tojson()) assert drop_fbits(from_ROOT["TH1"]) == drop_fbits( streamers["TH1"][7].tojson()) assert drop_fbits(from_ROOT["TNamed"]) == drop_fbits( streamers["TNamed"][1].tojson()) assert drop_fbits(from_ROOT["TObject"]) == drop_fbits( streamers["TObject"][1].tojson()) assert drop_fbits(from_ROOT["TAttLine"]) == drop_fbits( streamers["TAttLine"][2].tojson()) assert drop_fbits(from_ROOT["TAttFill"]) == drop_fbits( streamers["TAttFill"][2].tojson()) assert drop_fbits(from_ROOT["TAttMarker"]) == drop_fbits( streamers["TAttMarker"][2].tojson()) assert drop_fbits(from_ROOT["TAxis"]) == drop_fbits( streamers["TAxis"][10].tojson()) assert drop_fbits(from_ROOT["TAttAxis"]) == drop_fbits( streamers["TAttAxis"][4].tojson()) assert drop_fbits(from_ROOT["THashList"]) == drop_fbits( streamers["THashList"][0].tojson()) assert drop_fbits(from_ROOT["TList"]) == drop_fbits( streamers["TList"][5].tojson()) assert drop_fbits(from_ROOT["TSeqCollection"]) == drop_fbits( streamers["TSeqCollection"][0].tojson()) assert drop_fbits(from_ROOT["TCollection"]) == drop_fbits( streamers["TCollection"][3].tojson()) assert drop_fbits(from_ROOT["TString"]) == drop_fbits( streamers["TString"][2].tojson())
def test_map_string_int16(): with uproot4.open(skhep_testdata.data_path( "uproot-stl_containers.root"))["tree"] as tree: assert [ x.tolist() for x in tree["map_string_int16"].array(library="np") ] == [ { "one": 1 }, { "one": 1, "two": 2 }, { "one": 1, "two": 2, "three": 3 }, { "one": 1, "two": 2, "three": 3, "four": 4 }, { "one": 1, "two": 2, "three": 3, "four": 4, "five": 5 }, ]
def test_branchname(): with uproot.open( skhep_testdata.data_path("uproot-sample-6.20.04-uncompressed.root") )["sample"] as sample: assert sample.arrays("i4", library="np")["i4"].tolist() == list(range(-15, 15)) arrays = sample.arrays(["i4", "i8"], library="np") assert set(arrays.keys()) == set(["i4", "i8"]) assert arrays["i4"].tolist() == list(range(-15, 15)) assert arrays["i8"].tolist() == list(range(-15, 15)) arrays = sample.arrays(filter_name="/i[48]/", library="np") assert set(arrays.keys()) == set(["i4", "i8"]) assert arrays["i4"].tolist() == list(range(-15, 15)) assert arrays["i8"].tolist() == list(range(-15, 15)) arrays = sample.arrays(filter_name=["/i[12]/", "/i[48]/"], library="np") assert set(arrays.keys()) == set(["i1", "i2", "i4", "i8"]) assert arrays["i1"].tolist() == list(range(-15, 15)) assert arrays["i2"].tolist() == list(range(-15, 15)) assert arrays["i4"].tolist() == list(range(-15, 15)) assert arrays["i8"].tolist() == list(range(-15, 15)) arrays = sample.arrays(filter_name="i*", library="np") assert set(arrays.keys()) == set(["i1", "i2", "i4", "i8"]) assert arrays["i1"].tolist() == list(range(-15, 15)) assert arrays["i2"].tolist() == list(range(-15, 15)) assert arrays["i4"].tolist() == list(range(-15, 15)) assert arrays["i8"].tolist() == list(range(-15, 15)) arrays = sample.arrays(["i4", "i8"], filter_name="u*", library="np") assert set(arrays.keys()) == set(["i4", "i8"]) assert arrays["i4"].tolist() == list(range(-15, 15)) assert arrays["i8"].tolist() == list(range(-15, 15))
def test_map_string_set_string(): with uproot4.open(skhep_testdata.data_path( "uproot-stl_containers.root"))["tree"] as tree: assert [ x.tolist() for x in tree["map_string_set_string"].array(library="np") ] == [ { "one": set(["one"]) }, { "one": set(["one"]), "two": set(["one", "two"]) }, { "one": set(["one"]), "two": set(["one", "two"]), "three": set(["one", "two", "three"]), }, { "one": set(["one"]), "two": set(["one", "two"]), "three": set(["one", "two", "three"]), "four": set(["one", "two", "three", "four"]), }, { "one": set(["one"]), "two": set(["one", "two"]), "three": set(["one", "two", "three"]), "four": set(["one", "two", "three", "four"]), "five": set(["one", "two", "three", "four", "five"]), }, ]
def test_interpretation(): with uproot.open( skhep_testdata.data_path("uproot-sample-6.20.04-uncompressed.root") )["sample"] as sample: assert sample["i2"].array(">u2", library="np").tolist() == list( range(65521, 65536) ) + list(range(0, 15)) arrays = sample.arrays({"i1": ">u1", "i2": ">u2"}, library="np") assert set(arrays.keys()) == set(["i1", "i2"]) assert arrays["i1"].tolist() == list(range(241, 256)) + list(range(0, 15)) assert arrays["i2"].tolist() == list(range(65521, 65536)) + list(range(0, 15)) arrays = sample.arrays([("i1", ">u1"), ("i2", ">u2")], library="np", how=tuple) assert isinstance(arrays, tuple) and len(arrays) == 2 assert arrays[0].tolist() == list(range(241, 256)) + list(range(0, 15)) assert arrays[1].tolist() == list(range(65521, 65536)) + list(range(0, 15)) arrays = sample.arrays({"i1": ">u1", "i2": None}, library="np") assert set(arrays.keys()) == set(["i1", "i2"]) assert arrays["i1"].tolist() == list(range(241, 256)) + list(range(0, 15)) assert arrays["i2"].tolist() == list(range(-15, 15)) arrays = sample.arrays([("i1", ">u1"), ("i2", None)], library="np", how=tuple) assert isinstance(arrays, tuple) and len(arrays) == 2 assert arrays[0].tolist() == list(range(241, 256)) + list(range(0, 15)) assert arrays[1].tolist() == list(range(-15, 15)) with pytest.raises(ValueError): sample.arrays([("i1", ">u1"), ("i1", None)], library="np", how=tuple)
def test_map_string_tstring(): with uproot4.open(skhep_testdata.data_path( "uproot-stl_containers.root"))["tree"] as tree: assert [ x.tolist() for x in tree["map_string_tstring"].array(library="np") ] == [ { "one": "ONE" }, { "one": "ONE", "two": "TWO" }, { "one": "ONE", "two": "TWO", "three": "THREE" }, { "one": "ONE", "two": "TWO", "three": "THREE", "four": "FOUR" }, { "one": "ONE", "two": "TWO", "three": "THREE", "four": "FOUR", "five": "FIVE", }, ]
def test_compute(): with uproot.open( skhep_testdata.data_path("uproot-sample-6.20.04-uncompressed.root") )["sample"] as sample: assert sample.arrays("i4 + 100", library="np")["i4 + 100"].tolist() == list( range(85, 115) ) arrays = sample.arrays(["i4 + 100", "i8 + 100"], library="np") assert set(arrays.keys()) == set(["i4 + 100", "i8 + 100"]) assert arrays["i4 + 100"].tolist() == list(range(85, 115)) assert arrays["i8 + 100"].tolist() == list(range(85, 115)) arrays = sample.arrays(["i4 + 100", "i4 + 200"], library="np") assert set(arrays.keys()) == set(["i4 + 100", "i4 + 200"]) assert arrays["i4 + 100"].tolist() == list(range(85, 115)) assert arrays["i4 + 200"].tolist() == list(range(185, 215)) arrays = sample.arrays(["i4 + 100", "i4 + 100"], library="np") assert set(arrays.keys()) == set(["i4 + 100"]) assert arrays["i4 + 100"].tolist() == list(range(85, 115)) arrays = sample.arrays(["i4 + 100", "i4 + 100"], library="np", how=tuple) assert isinstance(arrays, tuple) and len(arrays) == 2 assert arrays[0].tolist() == list(range(85, 115)) assert arrays[1].tolist() == list(range(85, 115))