From 32f54ee5bfc1f054bef50e5b20c1d81c4c0537f0 Mon Sep 17 00:00:00 2001
From: Tamas Gal <tgal@km3net.de>
Date: Sun, 22 Dec 2019 14:49:40 +0100
Subject: [PATCH] Allow slicing
---
km3io/offline.py | 11 +++++++++--
tests/test_offline.py | 6 ++++--
2 files changed, 13 insertions(+), 4 deletions(-)
diff --git a/km3io/offline.py b/km3io/offline.py
index 6823a5d..ee868ef 100644
--- a/km3io/offline.py
+++ b/km3io/offline.py
@@ -570,8 +570,15 @@ class OfflineTracks:
setattr(self, k, v)
def __getitem__(self, item):
- return OfflineTrack(self._keys, [v[item] for v in self._values],
- fit_keys=self._fit_keys)
+ if isinstance(item, int):
+ return OfflineTrack(self._keys, [v[item] for v in self._values],
+ fit_keys=self._fit_keys)
+ else:
+ return OfflineTracks(
+ self._keys,
+ [v[item] for v in self._values],
+ fit_keys=self._fit_keys
+ )
def __len__(self):
try:
diff --git a/tests/test_offline.py b/tests/test_offline.py
index a95f5d7..9c73dbf 100644
--- a/tests/test_offline.py
+++ b/tests/test_offline.py
@@ -309,8 +309,10 @@ class TestOfflineTracks(unittest.TestCase):
def test_slicing(self):
tracks = self.tracks
assert 10 == len(tracks)
- track_selection = tracks[2:4]
- assert 3 == len(track_selection)
+ track_selection = tracks[2:7]
+ assert 5 == len(track_selection)
+ track_selection_2 = tracks[1:3]
+ assert 2 == len(track_selection_2)
class TestOfflineTrack(unittest.TestCase):
--
GitLab