Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions docs/guide/cli.md
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ That's why taskiq can auto-discover tasks in current directory recursively.
We have two options for this:

- `--tasks-pattern` or `-tp`.
It's a name of files to import. By default is searches for all `tasks.py` files.
It's a glob pattern of files to import. By default it is `**/tasks.py` which searches for all `tasks.py` files. May be specified multiple times.
- `--fs-discover` or `-fsd`. This option enables search of task files in current directory recursively, using the given pattern.

### Acknowledgements
Expand Down Expand Up @@ -118,7 +118,7 @@ taskiq scheduler my_project.broker:scheduler my_project.module1 my_project.modul
Path to scheduler is the only required argument.

- `--tasks-pattern` or `-tp`.
It's a name of files to import. By default is searches for all `tasks.py` files.
It's a glob pattern of files to import. By default it is `**/tasks.py` which searches for all `tasks.py` files. May be specified multiple times.
- `--fs-discover` or `-fsd`. This option enables search of task files in current directory recursively, using the given pattern.
- `--no-configure-logging` - use this parameter if your application configures custom logging.
- `--log-level` is used to set a log level (default `INFO`).
Expand Down
15 changes: 11 additions & 4 deletions taskiq/cli/scheduler/args.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ class SchedulerArgs:
log_level: str = LogLevel.INFO.name
configure_logging: bool = True
fs_discover: bool = False
tasks_pattern: str = "tasks.py"
tasks_pattern: Sequence[str] = ("**/tasks.py",)
skip_first_run: bool = False

@classmethod
Expand Down Expand Up @@ -52,8 +52,9 @@ def from_cli(cls, args: Optional[Sequence[str]] = None) -> "SchedulerArgs":
parser.add_argument(
"--tasks-pattern",
"-tp",
default="tasks.py",
help="Name of files in which taskiq will try to find modules.",
default=["**/tasks.py"],
action="append",
help="Glob patterns of files in which taskiq will try to find the tasks.",
)
parser.add_argument(
"--log-level",
Expand All @@ -76,4 +77,10 @@ def from_cli(cls, args: Optional[Sequence[str]] = None) -> "SchedulerArgs":
"This option skips running tasks immediately after scheduler start."
),
)
return cls(**parser.parse_args(args).__dict__)

namespace = parser.parse_args(args)
# If there are any patterns specified, remove default.
# This is an argparse limitation.
if len(namespace.tasks_pattern) > 1:
namespace.tasks_pattern.pop(0)
return cls(**namespace.__dict__)
23 changes: 16 additions & 7 deletions taskiq/cli/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from importlib import import_module
from logging import getLogger
from pathlib import Path
from typing import Any, Generator, List
from typing import Any, Generator, List, Sequence, Union

from taskiq.utils import remove_suffix

Expand Down Expand Up @@ -69,7 +69,11 @@ def import_from_modules(modules: List[str]) -> None:
logger.exception(err)


def import_tasks(modules: List[str], pattern: str, fs_discover: bool) -> None:
def import_tasks(
modules: List[str],
pattern: Union[str, Sequence[str]],
fs_discover: bool,
) -> None:
"""
Import tasks modules.

Expand All @@ -82,9 +86,14 @@ def import_tasks(modules: List[str], pattern: str, fs_discover: bool) -> None:
from filesystem.
"""
if fs_discover:
for path in Path().rglob(pattern):
modules.append(
remove_suffix(str(path), ".py").replace(os.path.sep, "."),
)

if isinstance(pattern, str):
pattern = (pattern,)
discovered_modules = set()
for glob_pattern in pattern:
for path in Path().glob(glob_pattern):
discovered_modules.add(
remove_suffix(str(path), ".py").replace(os.path.sep, "."),
)

modules.extend(list(discovered_modules))
import_from_modules(modules)
11 changes: 8 additions & 3 deletions taskiq/cli/worker/args.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ class WorkerArgs:

broker: str
modules: List[str]
tasks_pattern: str = "tasks.py"
tasks_pattern: Sequence[str] = ("**/tasks.py",)
fs_discover: bool = False
configure_logging: bool = True
log_level: LogLevel = LogLevel.INFO
Expand Down Expand Up @@ -87,8 +87,9 @@ def from_cli(
parser.add_argument(
"--tasks-pattern",
"-tp",
default="tasks.py",
help="Name of files in which taskiq will try to find modules.",
default=["**/tasks.py"],
action="append",
help="Glob patterns of files in which taskiq will try to find the tasks.",
)
parser.add_argument(
"modules",
Expand Down Expand Up @@ -198,4 +199,8 @@ def from_cli(
)

namespace = parser.parse_args(args)
# If there are any patterns specified, remove default.
# This is an argparse limitation.
if len(namespace.tasks_pattern) > 1:
namespace.tasks_pattern.pop(0)
return WorkerArgs(**namespace.__dict__)
43 changes: 43 additions & 0 deletions tests/cli/test_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
from unittest.mock import patch

from taskiq.cli.utils import import_tasks


def test_import_tasks_list_pattern() -> None:
modules = ["taskiq.tasks"]
with patch("taskiq.cli.utils.import_from_modules", autospec=True) as mock:
import_tasks(modules, ["tests/**/test_utils.py"], True)
assert set(modules) == {
"taskiq.tasks",
"tests.test_utils",
"tests.cli.test_utils",
}
mock.assert_called_with(modules)


def test_import_tasks_str_pattern() -> None:
modules = ["taskiq.tasks"]
with patch("taskiq.cli.utils.import_from_modules", autospec=True) as mock:
import_tasks(modules, "tests/**/test_utils.py", True)
assert set(modules) == {
"taskiq.tasks",
"tests.test_utils",
"tests.cli.test_utils",
}
mock.assert_called_with(modules)


def test_import_tasks_empty_pattern() -> None:
modules = ["taskiq.tasks"]
with patch("taskiq.cli.utils.import_from_modules", autospec=True) as mock:
import_tasks(modules, [], True)
assert modules == ["taskiq.tasks"]
mock.assert_called_with(modules)


def test_import_tasks_no_discover() -> None:
modules = ["taskiq.tasks"]
with patch("taskiq.cli.utils.import_from_modules", autospec=True) as mock:
import_tasks(modules, "tests/**/test_utils.py", False)
assert modules == ["taskiq.tasks"]
mock.assert_called_with(modules)