Skip to content

Commit f6c96b2

Browse files
authored
support discovering pyc/pyd/so modules (#391)
* support discovering pyc/pyd/so modules * fix lint error * add trailing comma
1 parent d7e1b79 commit f6c96b2

File tree

2 files changed

+38
-5
lines changed

2 files changed

+38
-5
lines changed

taskiq/cli/utils.py

+10-5
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,6 @@
66
from pathlib import Path
77
from typing import Any, Generator, List, Sequence, Union
88

9-
from taskiq.utils import remove_suffix
10-
119
logger = getLogger("taskiq.worker")
1210

1311

@@ -91,9 +89,16 @@ def import_tasks(
9189
discovered_modules = set()
9290
for glob_pattern in pattern:
9391
for path in Path().glob(glob_pattern):
94-
discovered_modules.add(
95-
remove_suffix(str(path), ".py").replace(os.path.sep, "."),
96-
)
92+
if path.is_file():
93+
if path.suffix in (".py", ".pyc", ".pyd", ".so"):
94+
# remove all suffixes
95+
prefix = path.name.partition(".")[0]
96+
discovered_modules.add(
97+
str(path.with_name(prefix)).replace(os.path.sep, "."),
98+
)
99+
# ignore other files
100+
else:
101+
discovered_modules.add(str(path).replace(os.path.sep, "."))
97102

98103
modules.extend(list(discovered_modules))
99104
import_from_modules(modules)

tests/cli/test_utils.py

+28
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from contextlib import suppress
2+
from pathlib import Path
13
from unittest.mock import patch
24

35
from taskiq.cli.utils import import_tasks
@@ -41,3 +43,29 @@ def test_import_tasks_no_discover() -> None:
4143
import_tasks(modules, "tests/**/test_utils.py", False)
4244
assert modules == ["taskiq.tasks"]
4345
mock.assert_called_with(modules)
46+
47+
48+
def test_import_tasks_non_py_list_pattern() -> None:
49+
modules = ["taskiq.tasks"]
50+
with patch("taskiq.cli.utils.import_from_modules", autospec=True) as mock:
51+
pathes = (
52+
Path("tests/test1.so"),
53+
Path("tests/cli/test2.cpython-313-darwin.so"),
54+
)
55+
for path in pathes:
56+
path.touch()
57+
58+
try:
59+
import_tasks(modules, ["tests/**/test_utils.py", "tests/**/*.so"], True)
60+
assert set(modules) == {
61+
"taskiq.tasks",
62+
"tests.test_utils",
63+
"tests.cli.test_utils",
64+
"tests.test1",
65+
"tests.cli.test2",
66+
}
67+
mock.assert_called_with(modules)
68+
finally:
69+
for path in pathes:
70+
with suppress(FileNotFoundError):
71+
path.unlink()

0 commit comments

Comments
 (0)