Add factories (#349)

* Add factories support

---------

Co-authored-by: ddmytruk <[email protected]>

Author: 2 people

taskiq/brokers/inmemory_broker.py

@@ -87,7 +87,7 @@ async def set_progress(
         progress: TaskProgress[Any],
     ) -> None:
         """
-        Set progress of task exection.
+        Set progress of task execution.
 
         :param task_id: task id
         :param progress: task execution progress

taskiq/cli/scheduler/args.py

@@ -32,7 +32,10 @@ def from_cli(cls, args: Optional[Sequence[str]] = None) -> "SchedulerArgs":
             formatter_class=ArgumentDefaultsHelpFormatter,
             description="Subcommand to run scheduler",
         )
-        parser.add_argument("scheduler", help="Path to scheduler")
+        parser.add_argument(
+            "scheduler",
+            help="Path to scheduler or scheduler factory function",
+        )
         parser.add_argument(
             "modules",
             help="List of modules where to look for tasks.",

taskiq/cli/scheduler/run.py

@@ -1,4 +1,5 @@
 import asyncio
+import inspect
 import sys
 from datetime import datetime, timedelta
 from logging import basicConfig, getLevelName, getLogger
@@ -192,15 +193,19 @@ async def run_scheduler(args: SchedulerArgs) -> None:
             ),
         )
     getLogger("taskiq").setLevel(level=getLevelName(args.log_level))
+
     if isinstance(args.scheduler, str):
         scheduler = import_object(args.scheduler)
+        if inspect.isfunction(scheduler):
+            scheduler = scheduler()
     else:
         scheduler = args.scheduler
     if not isinstance(scheduler, TaskiqScheduler):
         logger.error(
             "Imported scheduler is not a subclass of TaskiqScheduler.",
         )
         sys.exit(1)
+
     scheduler.broker.is_scheduler_process = True
     import_tasks(args.modules, args.tasks_pattern, args.fs_discover)
     for source in scheduler.sources:

taskiq/cli/worker/args.py

@@ -62,7 +62,7 @@ def from_cli(
         parser.add_argument(
             "broker",
             help=(
-                "Where to search for broker. "
+                "Where to search for broker or broker factory function. "
                 "This string must be specified in "
                 "'module.module:variable' format."
             ),

taskiq/cli/worker/run.py

@@ -1,4 +1,5 @@
 import asyncio
+import inspect
 import logging
 import os
 import signal
@@ -126,9 +127,16 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
     # broker is running as a worker.
     # We must set this field before importing tasks,
     # so broker will remember all tasks it's related to.
+
     broker = import_object(args.broker)
+    if inspect.isfunction(broker):
+        broker = broker()
     if not isinstance(broker, AsyncBroker):
-        raise ValueError("Unknown broker type. Please use AsyncBroker instance.")
+        raise ValueError(
+            "Unknown broker type. Please use AsyncBroker instance "
+            "or pass broker factory function that returns an AsyncBroker instance.",
+        )
+
     broker.is_worker_process = True
     import_tasks(args.modules, args.tasks_pattern, args.fs_discover)