Class TestIterator

dali/test/python/test_external_source_impl.py:113–150

class TestIterator:
    def __init__(self, n, batch_size, dims=[2], as_tensor=False):
        self.batch_size = batch_size
        self.dims = dims
        self.n = n
        self.as_tensor = as_tensor
        self.i = 0

    def __len__(self):
        return self.n

    def __iter__(self):
        # return a copy, so that the iteration number doesn't collide
        return TestIterator(self.n, self.batch_size, self.dims, self.as_tensor)

    def __next__(self):
        random_seed(12345 * self.i + 4321)

        def generate(dim):
            if self.as_tensor:
                shape = random_int(1, 10, [dim]).tolist()
                return random_array([self.batch_size] + shape)
            else:
                return [
                    random_array(random_int(1, 10, [dim]).tolist()) for _ in range(self.batch_size)
                ]

        if self.i < self.n:
            self.i += 1
            if isinstance(self.dims, (list, tuple)):
                return [generate(d) for d in self.dims]
            else:
                return generate(self.dims)
        else:
            self.i = 0
            raise StopIteration

    next = __next__


class SampleIterator: