loading/loaderiterator.py

from pathlib import Path
from typing import List, Dict, Optional

from loading.serialization import Serializer

class LoaderIterator:
    """Iterator that loads data from multiple files in batches."""

    def __init__(self,
                 serializer: Serializer,
                 num_files_per_iteration: int,
                 load_paths: Optional[List[Path]] = None) -> None:
        self.serializer = serializer
        self.num_files_per_iteration = num_files_per_iteration
        self._load_paths = load_paths
        self._current_iteration = None

    @property
    def load_paths(self) -> Optional[List[Path]]:
        return self._load_paths

    @load_paths.setter
    def load_paths(self, load_paths: List[Path]) -> None:
        self._load_paths = load_paths

    def __iter__(self):
        self._current_iteration = 0
        return self

    def __next__(self) -> List[Dict]:
        if self._did_load_all_batches():
            raise StopIteration
        data_batch = self._load_data_batch()
        self._current_iteration += 1
        return data_batch

    def _did_load_all_batches(self) -> bool:
        if self._current_iteration >= len(self._load_paths) / self.num_files_per_iteration:
            return True
        return False

    def _load_data_batch(self) -> List[Dict]:
        start_index = self._current_iteration * self.num_files_per_iteration
        stop_index = start_index + self.num_files_per_iteration
        return [self.serializer.load(load_path) for load_path in
                self._load_paths[start_index:stop_index] if load_path.exists()]