diff --git a/CHANGELOG.rst b/CHANGELOG.rst index baae35a535516f9e6a7e9d135c7154d518b094ad..cef33b1d8dafb8d3a9ffb8e8a4f7fa5081f402b6 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -4,6 +4,10 @@ Unreleased changes Version 0 --------- +0.6.2 / 2019-12-22 +~~~~~~~~~~~~~~~~~~ +* Fixes slicing of ``OfflineTracks`` + 0.6.1 / 2019-12-21 ~~~~~~~~~~~~~~~~~~ * Minor cleanup diff --git a/km3io/offline.py b/km3io/offline.py index 6823a5d1aabf514e02509f24a05fa8ab8f9c6d4e..ee868ef9402aa4d2487b1449b710df202aaad51f 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 2dbae94258c6f4aa3d9dff958066d26b11796ec8..cdd476e80706bc0e0f8057ed6709beaed020e13b 100644 --- a/tests/test_offline.py +++ b/tests/test_offline.py @@ -306,6 +306,25 @@ class TestOfflineTracks(unittest.TestCase): self.assertListEqual([0.230189, 0.230189, 0.218663], list(dir_z[0][:3])) + def test_slicing(self): + tracks = self.tracks + assert 10 == len(tracks) + track_selection = tracks[2:7] + assert 5 == len(track_selection) + track_selection_2 = tracks[1:3] + assert 2 == len(track_selection_2) + for _slice in [ + slice(0, 0), + slice(0, 1), + slice(0, 2), + slice(1, 5), + slice(3, -2) + ]: + self.assertListEqual( + list(tracks.E[:,0][_slice]), + list(tracks[_slice].E[:,0]) + ) + class TestOfflineTrack(unittest.TestCase): def setUp(self):