From 8a42e7de07a74a81e300410bfcc2b1c5b3339820 Mon Sep 17 00:00:00 2001 From: ioanaif Date: Thu, 15 Feb 2024 18:15:18 +0100 Subject: [PATCH] fix: allow colon in key names (#1127) * fix: allow column in key names * style: pre-commit fixes * Rename test file to match PR number. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- src/uproot/reading.py | 2 ++ .../test_1127_fix_allow_colon_in_key_names.py | 26 +++++++++++++++++++ 2 files changed, 28 insertions(+) create mode 100644 tests/test_1127_fix_allow_colon_in_key_names.py diff --git a/src/uproot/reading.py b/src/uproot/reading.py index d046116a8..20700f4d3 100644 --- a/src/uproot/reading.py +++ b/src/uproot/reading.py @@ -2060,6 +2060,8 @@ def __getitem__(self, where): keys=[key.fName for key in last._keys], file_path=self._file.file_path, ) + elif ":" in item and item in step: + return self.key(where).get() else: last = step step = step[item] diff --git a/tests/test_1127_fix_allow_colon_in_key_names.py b/tests/test_1127_fix_allow_colon_in_key_names.py new file mode 100644 index 000000000..698e45d67 --- /dev/null +++ b/tests/test_1127_fix_allow_colon_in_key_names.py @@ -0,0 +1,26 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE + +import awkward as ak +import uproot +import skhep_testdata +import numpy +import os + + +def test_colon_in_path_and_name(tmp_path): + newfile = os.path.join(tmp_path, "test_colon_in_name.root") + with uproot.recreate(newfile) as f: + f["one:two"] = "together" + array = ak.Array(["one", "two", "three"]) + f["one"] = {"two": array} + + with uproot.open(newfile) as f: + f["one:two"] == "together" + f["one"]["two"].array() == ["one", "two", "three"] + + +def test_colon_reading_in_path(): + with uproot.open( + skhep_testdata.data_path("uproot-small-evnt-tree-fullsplit.root") + ) as f: + f["tree:evt/P3/P3.Py"].array() == numpy.arange(100)