ibonn/dammy

from .exceptions import MaximumRetriesExceededException

class Iterator:
    """
    Iterators provide a way to generate all instances given a class by iterating over them.
    """
    def __init__(self, c, dataset, raw=False, localization=None):
        """
        Create a new iterator.

        :param c: The class for the generator
        :param dataset: The dataset from which all referenced fields will be retrieved
        :param raw: Determines wether the iterator generates raw values or not
        :type c: :class:`dammy.core.BaseGenerator` or any class inheriting from it
        :type dataset: :class:`dammy.db.DatasetGenerator` or dict
        :type raw: bool
        :returns: A Python iterator
        """
        self.generator = c()
        self.dataset = dataset
        self.raw = raw
        self.num_generated_instances = 0
        self.localization = localization

    def __iter__(self):
        # This method has been implemented to make the object iterable
        return self

    def __next__(self):
        """
        Get the next value.
        :returns: An instance generated by the associated generator
        """
        try:
            if self.raw:
                instance = self.generator.generate_raw(self.dataset, self.localization)
            else:
                instance = self.generator.generate(self.dataset, self.localization)

        except MaximumRetriesExceededException:
            raise StopIteration
        self.num_generated_instances += 1
        return instance