add: pipelines nodes can now spawn one to many jobs via yield

pipeline.py

@@ -27,24 +27,24 @@ class Node:
             job: Job = await self.input_queue.get()
             self._jobs_dequeued += 1
             if self.sequential_node == False:
-                await self.process_job(job)
-                if self.output_queue is not None:
-                    await self.output_queue.put(job)
-                if self.job_sync is not None:
-                    self.job_sync.append(job)
-                self._jobs_processed += 1
+                async for job in self.process_job(job):
+                    if self.output_queue is not None:
+                        await self.output_queue.put(job)
+                    if self.job_sync is not None:
+                        self.job_sync.append(job)
+                    self._jobs_processed += 1
             else:
                 # ensure that jobs are processed in order
                 self.buffer[job.id] = job
                 while self.next_i in self.buffer:
                     job = self.buffer.pop(self.next_i)
-                    await self.process_job(job)
+                    async for job in self.process_job(job):
+                        if self.output_queue is not None:
+                            await self.output_queue.put(job)
+                        if self.job_sync is not None:
+                            self.job_sync.append(job)
+                        self._jobs_processed += 1
                     self.next_i += 1
-                    if self.output_queue is not None:
-                        await self.output_queue.put(job)
-                    if self.job_sync is not None:
-                        self.job_sync.append(job)
-                    self._jobs_processed += 1
 
     async def process_job(self, job: Job):
         raise NotImplementedError()

tests/test_pipeline.py

@@ -12,6 +12,7 @@ from pipeline import Pipeline, Node, Job
 class Node1(Node):
     async def process_job(self, job: Job):
         job.data += f' (processed by node 1, worker {self.worker_id})'
+        yield job
 
 
 class Node2(Node):
@@ -19,12 +20,14 @@ class Node2(Node):
         sleep_duration = 0.08 + 0.04 * random.random()
         await asyncio.sleep(sleep_duration)
         job.data += f' (processed by node 2, worker {self.worker_id})'
+        yield job
 
 
 class Node3(Node):
     async def process_job(self, job: Job):
         job.data += f' (processed by node 3, worker {self.worker_id})'
         print(f'{job.id} - {job.data}')
+        yield job
 
 
 class TestPipeline(unittest.TestCase):
@@ -63,17 +66,17 @@ class TestPipeline(unittest.TestCase):
         asyncio.run(self._test_pipeline_edge_cases())
 
 
-    # def test_pipeline_keeps_order(self):
-    #     self.pipeline = Pipeline()
-    #     self.job_sync = []
-    #     num_jobs = 100
-    #     start_time = time.time()
-    #     asyncio.run(self._test_pipeline(num_jobs))
-    #     end_time = time.time()
-    #     print(f"Pipeline processed in {end_time - start_time} seconds.")
-    #     self.assertEqual(len(self.job_sync), num_jobs)
-    #     for i, job in enumerate(self.job_sync):
-    #         self.assertEqual(i, job.id)
+    def test_pipeline_keeps_order(self):
+        self.pipeline = Pipeline()
+        self.job_sync = []
+        num_jobs = 100
+        start_time = time.time()
+        asyncio.run(self._test_pipeline(num_jobs))
+        end_time = time.time()
+        print(f"Pipeline processed in {end_time - start_time} seconds.")
+        self.assertEqual(len(self.job_sync), num_jobs)
+        for i, job in enumerate(self.job_sync):
+            self.assertEqual(i, job.id)
 
 
 if __name__ == '__main__':