From fcca3704168e3d530abd6cfcbe063299b0a24225 Mon Sep 17 00:00:00 2001
From: Tamas Gal <tgal@km3net.de>
Date: Wed, 1 Apr 2020 16:26:28 +0200
Subject: [PATCH] Add iteration
---
km3io/tools.py | 18 +++++++++++++++++-
tests/test_offline.py | 10 ++++++++++
2 files changed, 27 insertions(+), 1 deletion(-)
diff --git a/km3io/tools.py b/km3io/tools.py
index 0cf10f9..e443357 100644
--- a/km3io/tools.py
+++ b/km3io/tools.py
@@ -55,6 +55,8 @@ class Branch:
self._subbranches = []
self._subbranchmaps = subbranchmaps
+ self._iterator_index = 0
+
if keymap is None:
self._initialise_keys() #
else:
@@ -125,8 +127,22 @@ class Branch:
self._branch[self._keymap['id']].lazyarray(
basketcache=BASKET_CACHE), self._index_chain))
+ def __iter__(self):
+ self._iterator_index = 0
+ return self
+
+ def __next__(self):
+ idx = self._iterator_index
+ self._iterator_index += 1
+ if idx >= len(self):
+ raise StopIteration
+ return self[idx]
+
def __str__(self):
- return "Number of elements: {}".format(len(self._branch))
+ length = len(self)
+ return "{} ({}) with {} element{}".format(self.__class__.__name__,
+ self._mapper.name, length,
+ 's' if length > 1 else '')
def __repr__(self):
length = len(self)
diff --git a/tests/test_offline.py b/tests/test_offline.py
index dfd8e7f..2dd1ee7 100644
--- a/tests/test_offline.py
+++ b/tests/test_offline.py
@@ -161,6 +161,16 @@ class TestOfflineEvents(unittest.TestCase):
assert np.allclose(self.events.hits[3:5][1][4].dom_id,
self.events[3:5][1][4].hits.dom_id)
+ def test_iteration(self):
+ i = 0
+ for event in self.events:
+ i += 1
+ assert 10 == i
+
+ def test_iteration_2(self):
+ n_hits = [e.n_hits for e in self.events]
+ assert np.allclose(n_hits, self.events.n_hits)
+
def test_str(self):
assert str(self.n_events) in str(self.events)
--
GitLab