From 077ef9bc0d56535f9e1432621a8d47eab03d454b Mon Sep 17 00:00:00 2001 From: Yousef Moazzam Date: Wed, 23 Apr 2025 10:05:47 +0100 Subject: [PATCH] Simplify `BlockSplitter`'s iterator implementation --- httomo/runner/block_split.py | 26 +++++++++----------------- 1 file changed, 9 insertions(+), 17 deletions(-) diff --git a/httomo/runner/block_split.py b/httomo/runner/block_split.py index f493fd7d7..e1e8aeec4 100644 --- a/httomo/runner/block_split.py +++ b/httomo/runner/block_split.py @@ -1,5 +1,4 @@ import math -from typing import Iterator from httomo.runner.dataset import DataSetBlock from httomo.runner.dataset_store_interfaces import DataSetSource @@ -22,6 +21,7 @@ def __init__(self, source: DataSetSource, max_slices: int): self._chunk_size = source.chunk_shape[source.slicing_dim] self._max_slices = int(min(max_slices, self._chunk_size)) self._num_blocks = math.ceil(self._chunk_size / self._max_slices) + self._current = 0 assert self._source.slicing_dim in [ 0, 1, @@ -41,20 +41,12 @@ def __getitem__(self, idx: int) -> DataSetBlock: len = min(self.slices_per_block, self._chunk_size - start) return self._source.read_block(start, len) - def __iter__(self) -> Iterator[DataSetBlock]: - class BlockIterator: - def __init__(self, splitter): - self.splitter = splitter - self._current = 0 + def __iter__(self): + return self - def __iter__(self) -> "BlockIterator": - return self # pragma: no cover - - def __next__(self) -> DataSetBlock: - if self._current >= len(self.splitter): - raise StopIteration - v = self.splitter[self._current] - self._current += 1 - return v - - return BlockIterator(self) + def __next__(self) -> DataSetBlock: + if self._current >= len(self): + raise StopIteration + v = self[self._current] + self._current += 1 + return v