DiamondLightSource / httomo

High-throughput tomography pipeline
https://diamondlightsource.github.io/httomo/
Other
5 stars 4 forks source link

Refactor `BlockSplitter`'s implementation of `__iter__()` to be clearer #336

Open yousefmoazzam opened 4 months ago

yousefmoazzam commented 4 months ago

The implementation is here: https://github.com/DiamondLightSource/httomo/blob/bce8fc7ffd411d22a07dc21c0b4cc090b4dcaa77/httomo/runner/block_split.py#L46-L60

It works as intended, but looks a bit complicated due to the definition of the BlockIterator class within the __iter__() method.

A cursory look at how to implement __iter__() (for example, https://realpython.com/python-iterators-iterables/#yielding-the-original-data) shows that this could be done without defining an extra class, doing something like:

def __iter__(self):
    return self

def __next__(self) -> DataSetBlock:
    if self._current >= len(self):
        raise StopIteration
    v = self[self._current]
    self._current += 1
    return v

and also making _current an attribute in BlockSplitter__init__():

def __init__(self, source: DataSetSource, max_slices: int):
    ...
    self._current = 0