From 7fb23f0ed6385de4fe97a5b388b18ad0f3edad60 Mon Sep 17 00:00:00 2001 From: zineb aly <aly.zineb.az@gmail.com> Date: Tue, 26 May 2020 09:28:59 +0200 Subject: [PATCH] monkey patch chunckedArray gititem --- km3io/__init__.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/km3io/__init__.py b/km3io/__init__.py index 1d7af38..2cb507f 100644 --- a/km3io/__init__.py +++ b/km3io/__init__.py @@ -1,3 +1,7 @@ +import numpy as np +import awkward as ak +import awkward1 as ak1 + from pkg_resources import get_distribution, DistributionNotFound version = get_distribution(__name__).version @@ -5,3 +9,18 @@ version = get_distribution(__name__).version from .offline import OfflineReader from .online import OnlineReader from .gseagen import GSGReader + +# to avoid infinite recursion +old_getitem = ak.ChunkedArray.__getitem__ + + +def new_getitem(self, item): + """Monkey patch the getitem in awkward.ChunkedArray to apply + awkward1.Array masks on awkward.ChunkedArray""" + if isinstance(item, (ak1.Array, ak.ChunkedArray)): + return ak1.Array(self)[item] + else: + return old_getitem(self, item) + + +ak.ChunkedArray.__getitem__ = new_getitem -- GitLab