Function test_external_source

dali/test/python/test_external_source_impl.py:505–556

()

def test_external_source():
    class TestIterator:
        def __init__(self, n):
            self.n = n

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

        def __next__(self):
            batch_1 = []
            batch_2 = []
            if self.i < self.n:
                batch_1.append(datapy.arange(0, 1, dtype=datapy.float32))
                batch_2.append(datapy.arange(0, 1, dtype=datapy.float32))
                self.i += 1
                return batch_1, batch_2
            else:
                self.i = 0
                raise StopIteration

        next = __next__

    class IterSetupPipeline(Pipeline):
        def __init__(self, iterator, num_threads, device_id):
            super().__init__(1, num_threads, device_id)
            self.input_1 = ops.ExternalSource()
            self.input_2 = ops.ExternalSource()
            self.iterator = iterator

        def define_graph(self):
            self.batch_1 = self.input_1()
            self.batch_2 = self.input_2()
            return [self.batch_1, self.batch_2]

        def iter_setup(self):
            batch_1, batch_2 = next(self.iterator)
            self.feed_input(self.batch_1, batch_1)
            self.feed_input(self.batch_2, batch_2)

    iter_num = 5
    iterator = iter(TestIterator(iter_num))
    pipe = IterSetupPipeline(iterator, 3, 0)

    i = 0
    while True:
        try:
            pipe.run()
            i += 1
        except StopIteration:
            break
    assert iter_num == i


def test_external_source_fail_missing_output():