From b45bc506e8457cce6aa5b4b85282c6f7823b0446 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Thu, 19 Nov 2020 11:41:13 -0600 Subject: [PATCH] Fixes lost cuts in iterate. (#195) * Fixes lost cuts in iterate. * Black and Flake8. --- tests/test_0194-fix-lost-cuts-in-iterate.py | 17 +++++++++++++++++ uproot4/behaviors/TBranch.py | 4 ++-- 2 files changed, 19 insertions(+), 2 deletions(-) create mode 100644 tests/test_0194-fix-lost-cuts-in-iterate.py diff --git a/tests/test_0194-fix-lost-cuts-in-iterate.py b/tests/test_0194-fix-lost-cuts-in-iterate.py new file mode 100644 index 000000000..b83f64b23 --- /dev/null +++ b/tests/test_0194-fix-lost-cuts-in-iterate.py @@ -0,0 +1,17 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/uproot4/blob/master/LICENSE + +from __future__ import absolute_import + +import numpy +import pytest +import skhep_testdata + +import uproot4 + + +def test(): + with uproot4.open(skhep_testdata.data_path("uproot-Zmumu.root")) as f: + for arrays in f["events"].iterate( + "px1", step_size=1000, cut="px1 > 0", library="np" + ): + assert numpy.all(arrays["px1"] > 0) diff --git a/uproot4/behaviors/TBranch.py b/uproot4/behaviors/TBranch.py index 56dc90e32..789094852 100644 --- a/uproot4/behaviors/TBranch.py +++ b/uproot4/behaviors/TBranch.py @@ -1375,13 +1375,13 @@ def iterate( self.object_path, ) - expression_context = [ + minimized_expression_context = [ (e, c) for e, c in expression_context if c["is_primary"] and not c["is_cut"] ] - arrays = library.group(output, expression_context, how) + arrays = library.group(output, minimized_expression_context, how) next_baskets = {} for branch, basket_num, basket in ranges_or_baskets: