Add --reload-dir argument (#459)

Author: montaguegabe

taskiq/cli/watcher.py

@@ -14,12 +14,13 @@ class FileWatcher:  # pragma: no cover
     def __init__(
         self,
         callback: Callable[..., None],
+        path: Path,
         use_gitignore: bool = True,
         **callback_kwargs: Any,
     ) -> None:
         self.callback = callback
         self.gitignore = None
-        gpath = Path("./.gitignore")
+        gpath = path / ".gitignore"
         if use_gitignore and gpath.exists():
             self.gitignore = parse_gitignore(gpath)
         self.callback_kwargs = callback_kwargs

taskiq/cli/worker/args.py

@@ -39,6 +39,7 @@ class WorkerArgs:
     no_parse: bool = False
     shutdown_timeout: float = 5
     reload: bool = False
+    reload_dirs: List[str] = field(default_factory=list)
     no_gitignore: bool = False
     max_async_tasks: int = 100
     receiver: str = "taskiq.receiver:Receiver"
@@ -172,6 +173,16 @@ def from_cli(
             help="Reload workers if file is changed. "
             "`reload` extra is required for this option.",
         )
+        parser.add_argument(
+            "--reload-dir",
+            action="append",
+            dest="reload_dirs",
+            default=[],
+            help=(
+                "Specify a directory to watch for changes. Can be specified "
+                "multiple times. Defaults to the current working directory."
+            ),
+        )
         parser.add_argument(
             "--do-not-use-gitignore",
             action="store_true",

taskiq/cli/worker/process_manager.py

@@ -6,6 +6,7 @@
 from dataclasses import dataclass
 from multiprocessing import Event, Process, Queue, current_process
 from multiprocessing.synchronize import Event as EventType
+from pathlib import Path
 from time import sleep
 from typing import Any, Callable, List, Optional
 
@@ -163,15 +164,19 @@ def __init__(
         self.action_queue: "Queue[ProcessActionBase]" = Queue(-1)
         self.args = args
         if args.reload and observer is not None:
-            observer.schedule(
-                FileWatcher(
-                    callback=schedule_workers_reload,
-                    use_gitignore=not args.no_gitignore,
-                    action_queue=self.action_queue,
-                ),
-                path=".",
-                recursive=True,
-            )
+            watch_paths = args.reload_dirs if args.reload_dirs else ["."]
+            for path_to_watch in watch_paths:
+                logger.debug(f"Watching directory: {path_to_watch}")
+                observer.schedule(
+                    FileWatcher(
+                        callback=schedule_workers_reload,
+                        path=Path(path_to_watch),
+                        use_gitignore=not args.no_gitignore,
+                        action_queue=self.action_queue,
+                    ),
+                    path=path_to_watch,
+                    recursive=True,
+                )
 
         shutdown_handler = get_signal_handler(self.action_queue, ShutdownAction())
         signal.signal(signal.SIGINT, shutdown_handler)