1. import asyncio
  2. import functools
  3. import inspect
  4. from typing import Callable, Any, Dict, List, Optional
  7. class TaskQueue:
  8. def __init__(self, name: str):
  9. self.name = name
  10. self._tasks: List[Any] = []
  11. self._lock = asyncio.Lock()
  13. async def enqueue(self, task: Callable, *args, **kwargs):
  14. async with self._lock:
  15. self._tasks.append((task, args, kwargs))
  16. await self._process()
  18. async def _process(self):
  19. if self._tasks:
  20. task, args, kwargs = self._tasks.pop(0)
  21. try:
  22. await task(*args, **kwargs)
  23. except Exception as e:
  24. print(f"Task {self.name} failed: {e}")
  25. finally:
  26. self._process() # Process next task
  29. class NestedTaskQueue:
  30. def __init__(self, name: str, parent: Optional['NestedTaskQueue'] = None):
  31. self.name = name
  32. self.parent = parent
  33. self._tasks: List[Any] = []
  34. self._lock = asyncio.Lock()
  36. async def enqueue(self, task: Callable, *args, **kwargs):
  37. async with self._lock:
  38. self._tasks.append((task, args, kwargs))
  39. await self._process()
  41. async def _process(self):
  42. if self._tasks:
  43. task, args, kwargs = self._tasks.pop(0)
  44. try:
  45. await task(*args, **kwargs)
  46. except Exception as e:
  47. print(f"Task {self.name} failed: {e}")
  48. finally:
  49. self._process()
  51. def get_child(self, name: str) -> 'NestedTaskQueue':
  52. if self.parent:
  53. return self.parent.get_child(name)
  54. else:
  55. return NestedTaskQueue(name, parent=self)
  58. def nest_task_queues(*args, **kwargs):
  59. """
  60. Creates a nested structure of task queues with optional flags.
  61. """
  63. root_queue = TaskQueue("root")
  64. queues = []
  66. for arg in args:
  67. if isinstance(arg, str):
  68. # Create a nested queue based on the string argument
  69. queue = TaskQueue(arg)
  70. queues.append(queue)
  71. elif isinstance(arg, NestedTaskQueue):
  72. queues.append(arg)
  73. else:
  74. raise TypeError("Each argument must be a string or a NestedTaskQueue instance.")
  76. # Connect the queues (optional, but allows for hierarchical task execution)
  77. for i in range(len(queues) - 1):
  78. queues[i].parent = queues[i+1]
  80. return root_queue, queues
  83. async def run_tasks(root_queue: TaskQueue, queues: List[NestedTaskQueue]):
  84. """
  85. Starts the task queues.
  86. """
  87. tasks = []
  89. # Enqueue tasks into the root queue
  90. for i in range(5):
  91. async def task_function(task_id):
  92. print(f"Running root task {task_id}")
  93. await asyncio.sleep(1)
  94. print(f"Root task {task_id} finished")
  95. tasks.append(root_queue.enqueue(task_function, i))
  97. # Enqueue tasks into child queues
  98. for i, queue in enumerate(queues):
  99. for j in range(3):
  100. async def child_task(task_id, queue_id):
  101. print(f"Running child task {task_id} in queue {queue_id}")
  102. await asyncio.sleep(0.5)
  103. print(f"Child task {task_id} in queue {queue_id} finished")
  104. queue.enqueue(child_task, j, queue_id)
  106. await asyncio.gather(*tasks)
  110. if __name__ == '__main__':
  111. # Example usage:
  112. root, nested_queues = nest_task_queues(
  113. "queue1",
  114. "queue2",
  115. NestedTaskQueue("queue3"),
  116. "queue4"
  117. )
  119. asyncio.run(run_tasks(root, nested_queues))

Add your comment