From 34346f84839d904138763bb82b28696f04357964 Mon Sep 17 00:00:00 2001
From: Tamas Gal <tgal@km3net.de>
Date: Wed, 1 Apr 2020 14:51:08 +0200
Subject: [PATCH] Add unfold indices routine
---
km3io/tools.py | 12 ++++++++++++
tests/test_tools.py | 19 ++++++++++++++++++-
2 files changed, 30 insertions(+), 1 deletion(-)
diff --git a/km3io/tools.py b/km3io/tools.py
index 641f4c4..418aafb 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 bceb4a0..3c4b655 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)
--
GitLab