|
| 1 | +# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE |
| 2 | + |
| 3 | +import os |
| 4 | + |
| 5 | +import pytest |
| 6 | + |
| 7 | +import uproot |
| 8 | + |
| 9 | +ak = pytest.importorskip("awkward") |
| 10 | + |
| 11 | +data = ak.Array( |
| 12 | + { |
| 13 | + "field1": [[(0, 1), (2, 3)], [], [(4, 5)]], |
| 14 | + "field2": [[{"x": 0, "y": 1}, {"x": 2, "y": 3}], [], [{"x": 4, "y": 5}]], |
| 15 | + "field3": [ |
| 16 | + {"x": [1, 2, 3], "y": [(1, 2), (2, 3), (3, 4)]}, |
| 17 | + {"x": [], "y": [(7, 8)]}, |
| 18 | + {"x": [9, 10], "y": [(11, 12), (13, 14)]}, |
| 19 | + ], |
| 20 | + "field4": [(1, [(1, 2), (3, 4)]), (2, []), (3, [(5, 6)])], |
| 21 | + "field5": [ |
| 22 | + { |
| 23 | + "x": [ |
| 24 | + {"up": [(0, 2), (1, 2)], "down": [[1, 2, 3], []]}, |
| 25 | + {"up": [], "down": [[4]]}, |
| 26 | + ], |
| 27 | + "y": [ |
| 28 | + ( |
| 29 | + {"left": [[0, "hi", 2.3]], "right": 6}, |
| 30 | + {"left": [[], [""]], "right": 8.0}, |
| 31 | + ) |
| 32 | + ], |
| 33 | + }, |
| 34 | + { |
| 35 | + "x": [ |
| 36 | + {"up": [(10, 2), (12, 2)], "down": []}, |
| 37 | + {"up": [(1, 2)], "down": [[], [4, 2, 3]]}, |
| 38 | + ], |
| 39 | + "y": [ |
| 40 | + ( |
| 41 | + {"left": [[23, 4.1, "hello"]], "right": 14}, |
| 42 | + {"left": [], "right": 16.0}, |
| 43 | + ) |
| 44 | + ], |
| 45 | + }, |
| 46 | + { |
| 47 | + "x": [ |
| 48 | + {"up": [], "down": [[4, 5], [1, 2, 3], []]}, |
| 49 | + {"up": [(0, 2), (1, 2)], "down": []}, |
| 50 | + ], |
| 51 | + "y": [ |
| 52 | + ( |
| 53 | + {"left": [[]], "right": 14}, |
| 54 | + {"left": [[2, 3], ["bye", ""]], "right": 16.0}, |
| 55 | + ) |
| 56 | + ], |
| 57 | + }, |
| 58 | + ], |
| 59 | + "field6": [(1, 2, 3, 4), (2, 5, 6, 7), ("hello", 2.3, 8, "")], |
| 60 | + } |
| 61 | +) |
| 62 | + |
| 63 | + |
| 64 | +def test_jagged_subfields(tmp_path): |
| 65 | + filepath = os.path.join(tmp_path, "test.root") |
| 66 | + |
| 67 | + with uproot.recreate(filepath) as file: |
| 68 | + obj = file.mkrntuple("ntuple", data) |
| 69 | + |
| 70 | + obj = uproot.open(filepath)["ntuple"] |
| 71 | + |
| 72 | + assert ak.array_equal(obj["field1"].array(), data["field1"]) |
| 73 | + assert ak.array_equal(obj["field1.0"].array(), data["field1"]["0"]) |
| 74 | + assert ak.array_equal(obj["field1.1"].array(), data["field1"]["1"]) |
| 75 | + |
| 76 | + assert ak.array_equal(obj["field2"].array(), data["field2"]) |
| 77 | + assert ak.array_equal(obj["field2.x"].array(), data["field2"]["x"]) |
| 78 | + assert ak.array_equal(obj["field2.y"].array(), data["field2"]["y"]) |
| 79 | + |
| 80 | + assert ak.array_equal(obj["field3"].array(), data["field3"]) |
| 81 | + assert ak.array_equal(obj["field3.x"].array(), data["field3"]["x"]) |
| 82 | + assert ak.array_equal(obj["field3.y"].array(), data["field3"]["y"]) |
| 83 | + assert ak.array_equal(obj["field3.y.0"].array(), data["field3"]["y"]["0"]) |
| 84 | + assert ak.array_equal(obj["field3.y.1"].array(), data["field3"]["y"]["1"]) |
| 85 | + |
| 86 | + assert ak.array_equal(obj["field4"].array(), data["field4"]) |
| 87 | + assert ak.array_equal(obj["field4.0"].array(), data["field4"]["0"]) |
| 88 | + assert ak.array_equal(obj["field4.1"].array(), data["field4"]["1"]) |
| 89 | + assert ak.array_equal(obj["field4.1.0"].array(), data["field4"]["1"]["0"]) |
| 90 | + assert ak.array_equal(obj["field4.1.1"].array(), data["field4"]["1"]["1"]) |
| 91 | + |
| 92 | + assert obj["field5"].array().tolist() == data["field5"].tolist() |
| 93 | + assert obj["field5.x"].array().tolist() == data["field5"]["x"].tolist() |
| 94 | + assert obj["field5.x.up"].array().tolist() == data["field5"]["x"]["up"].tolist() |
| 95 | + assert obj["field5.x.down"].array().tolist() == data["field5"]["x"]["down"].tolist() |
| 96 | + assert obj["field5.y"].array().tolist() == data["field5"]["y"].tolist() |
| 97 | + assert obj["field5.y.0"].array().tolist() == data["field5"]["y"]["0"].tolist() |
| 98 | + assert ( |
| 99 | + obj["field5.y.0.left"].array().tolist() |
| 100 | + == data["field5"]["y"]["0"]["left"].tolist() |
| 101 | + ) |
| 102 | + assert ( |
| 103 | + obj["field5.y.0.right"].array().tolist() |
| 104 | + == data["field5"]["y"]["0"]["right"].tolist() |
| 105 | + ) |
| 106 | + assert obj["field5.y.1"].array().tolist() == data["field5"]["y"]["1"].tolist() |
| 107 | + assert ( |
| 108 | + obj["field5.y.1.left"].array().tolist() |
| 109 | + == data["field5"]["y"]["1"]["left"].tolist() |
| 110 | + ) |
| 111 | + assert ( |
| 112 | + obj["field5.y.1.right"].array().tolist() |
| 113 | + == data["field5"]["y"]["1"]["right"].tolist() |
| 114 | + ) |
| 115 | + |
| 116 | + assert obj["field6"].array().tolist() == data["field6"].tolist() |
| 117 | + assert obj["field6.0"].array().tolist() == data["field6"]["0"].tolist() |
| 118 | + assert obj["field6.1"].array().tolist() == data["field6"]["1"].tolist() |
| 119 | + assert obj["field6.2"].array().tolist() == data["field6"]["2"].tolist() |
| 120 | + assert obj["field6.3"].array().tolist() == data["field6"]["3"].tolist() |
0 commit comments