diff --git a/km3io/tools.py b/km3io/tools.py index 641f4c4353fb0e1226591b3700d96a8b3455b965..418aafb4cd1f7fa54c9bdf411729eaeef7d72e7b 100644 --- a/km3io/tools.py +++ b/km3io/tools.py @@ -19,6 +19,18 @@ class cached_property: return prop +def _unfold_indices(obj, indices): + """Unfolds an index chain and returns the corresponding item""" + for depth, idx in enumerate(indices): + try: + obj = obj[idx] + except IndexError: + print("IndexError while accessing item '{}' at depth {} ({}) of " + "the index chain {}".format(repr(obj), depth, idx, indices)) + raise + return obj + + BranchMapper = namedtuple( "BranchMapper", ['name', 'key', 'extra', 'exclude', 'update', 'attrparser', 'flat']) diff --git a/tests/test_tools.py b/tests/test_tools.py index bceb4a0b45fb77784274513b337d83e79c7de303..3c4b6559e0158af6414111b94ca3cce717f2b6bb 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 import unittest -from km3io.tools import _to_num, cached_property +from km3io.tools import _to_num, cached_property, _unfold_indices class TestToNum(unittest.TestCase): @@ -20,3 +20,20 @@ class TestCachedProperty(unittest.TestCase): pass self.assertTrue(isinstance(Test.prop, cached_property)) + + +class TestUnfoldIndices(unittest.TestCase): + def test_unfold_indices(self): + data = range(10) + + indices = [slice(2, 5), 0] + assert data[indices[0]][indices[1]] == _unfold_indices(data, indices) + + indices = [slice(1, 9, 2), slice(1, 4), 2] + assert data[indices[0]][indices[1]][indices[2]] == _unfold_indices(data, indices) + + def test_unfold_indices_raises_index_error(self): + data = range(10) + indices = [slice(2, 5), 99] + with self.assertRaises(IndexError): + _unfold_indices(data, indices)