support discovering pyc/pyd/so modules

Author: keakon

taskiq/cli/utils.py

@@ -6,8 +6,6 @@
 from pathlib import Path
 from typing import Any, Generator, List, Sequence, Union
 
-from taskiq.utils import remove_suffix
-
 logger = getLogger("taskiq.worker")
 
 
@@ -91,9 +89,16 @@ def import_tasks(
         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, "."),
-                )
+                if path.is_file():
+                    if path.suffix in (".py", ".pyc", ".pyd", ".so"):
+                        # remove all suffixes
+                        prefix = path.name.partition(".")[0]
+                        discovered_modules.add(
+                            str(path.with_name(prefix)).replace(os.path.sep, ".")
+                        )
+                    # ignore other files
+                else:
+                    discovered_modules.add(str(path).replace(os.path.sep, "."))
 
         modules.extend(list(discovered_modules))
     import_from_modules(modules)

tests/cli/test_utils.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 from unittest.mock import patch
 
 from taskiq.cli.utils import import_tasks
@@ -41,3 +42,31 @@ def test_import_tasks_no_discover() -> None:
         import_tasks(modules, "tests/**/test_utils.py", False)
         assert modules == ["taskiq.tasks"]
         mock.assert_called_with(modules)
+
+
+def test_import_tasks_non_py_list_pattern() -> None:
+    modules = ["taskiq.tasks"]
+    with patch("taskiq.cli.utils.import_from_modules", autospec=True) as mock:
+        pathes = (
+            Path("tests/test1.so"),
+            Path("tests/cli/test2.cpython-313-darwin.so"),
+        )
+        for path in pathes:
+            path.touch()
+
+        try:
+            import_tasks(modules, ["tests/**/test_utils.py", "tests/**/*.so"], True)
+            assert set(modules) == {
+                "taskiq.tasks",
+                "tests.test_utils",
+                "tests.cli.test_utils",
+                "tests.test1",
+                "tests.cli.test2",
+            }
+            mock.assert_called_with(modules)
+        finally:
+            for path in pathes:
+                try:
+                    path.unlink()
+                except FileNotFoundError:
+                    pass