-
Notifications
You must be signed in to change notification settings - Fork 76
/
Copy pathtest_1272_basic_functionality.py
133 lines (97 loc) · 4.3 KB
/
test_1272_basic_functionality.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE
from __future__ import annotations
import pytest
from utils import run_test_in_pyodide
# Taken from test_0034_generic_objects_in_ttrees.py
@run_test_in_pyodide(test_file="uproot-HZZ-objects.root", packages=["pytest", "xxhash"])
def test_read_ttree(selenium):
import pytest
import uproot
awkward = pytest.importorskip("awkward")
with uproot.open("uproot-HZZ-objects.root")["events"] as tree:
result = tree["muonp4"].array(library="ak")
assert (
str(awkward.type(result))
== "2421 * var * TLorentzVector[fP: TVector3[fX: float64, "
"fY: float64, fZ: float64], fE: float64]"
)
assert result[0, 0, "fE"] == 54.77949905395508
assert result[0, 0, "fP", "fX"] == -52.89945602416992
assert result[0, 0, "fP", "fY"] == -11.654671669006348
assert result[0, 0, "fP", "fZ"] == -8.16079330444336
# Taken from test_0406_write_a_tree.py
@run_test_in_pyodide()
def test_write_ttree(selenium):
import numpy as np
import uproot
newfile = "newfile.root"
b1 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
b2 = [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9]
with uproot.recreate(newfile, compression=None) as fout:
tree = fout.mktree("t", {"b1": np.int32, "b2": np.float64}, "title")
assert tree._cascading._basket_capacity == 10
for _ in range(5):
fout["t"].extend({"b1": b1, "b2": b2})
assert tree._cascading._basket_capacity == 10
for _ in range(10):
fout["t"].extend({"b1": b1, "b2": b2})
assert tree._cascading._basket_capacity == 100
for _ in range(90):
fout["t"].extend({"b1": b1, "b2": b2})
assert tree._cascading._basket_capacity == 1000
with uproot.open(newfile) as fin:
assert fin.keys() == ["t;1"] # same cycle number
t2 = fin["t"]
assert t2.num_entries == len(b1) * 105
assert t2["b1"].array(library="np").tolist() == b1 * 105
assert t2["b2"].array(library="np").tolist() == b2 * 105
# Taken from test_1191_rntuple_fixes.py
@pytest.mark.skip(reason="Skipping until test files are available with RNTuple v1.0")
@run_test_in_pyodide(test_file="test_ntuple_extension_columns.root")
def test_read_rntuple(selenium):
import uproot
with uproot.open("test_ntuple_extension_columns.root") as f:
obj = f["EventData"]
assert len(obj.column_records) > len(obj.header.column_records)
assert len(obj.column_records) == 936
assert obj.column_records[903].first_element_index == 36
arrays = obj.arrays()
pbs = arrays[
"HLT_AntiKt4EMPFlowJets_subresjesgscIS_ftf_TLAAux::fastDIPS20211215_pb"
]
assert len(pbs) == 40
assert all(len(a) == 0 for a in pbs[:36])
assert next(i for i, a in enumerate(pbs) if len(a) != 0) == 36
jets = arrays["HLT_AntiKt4EMPFlowJets_subresjesgscIS_ftf_TLAAux:"]
assert len(jets.pt) == len(pbs)
# Taken from test_0034_generic_objects_in_ttrees.py
@pytest.mark.skip(reason="Skipping until test files are available with RNTuple v1.0")
@pytest.mark.network
@run_test_in_pyodide(packages=["requests"])
def test_read_ttree_http(selenium):
import awkward
import uproot
with uproot.open(
"https://github.com/scikit-hep/scikit-hep-testdata/raw/main/src/skhep_testdata/data/uproot-stl_containers.root",
handler=uproot.source.http.HTTPSource,
)["tree"] as tree:
assert awkward.to_list(tree["vector_int32"].array(library="ak")) == [
[1],
[1, 2],
[1, 2, 3],
[1, 2, 3, 4],
[1, 2, 3, 4, 5],
]
# Taken from test_1191_rntuple_fixes.py
@pytest.mark.skip(reason="Skipping until test files are available with RNTuple v1.0")
@pytest.mark.network
@run_test_in_pyodide(packages=["requests"])
def test_read_rntuple_http(selenium):
import uproot
with uproot.open(
"https://github.com/scikit-hep/scikit-hep-testdata/raw/main/src/skhep_testdata/data/Run2012BC_DoubleMuParked_Muons_rntuple_1000evts.root",
handler=uproot.source.http.HTTPSource,
) as f:
obj = f["Events"]
arrays = obj.arrays()
assert arrays["nMuon"].tolist() == [len(a) for a in arrays["Muon_pt"]]