From 038ef8aaf9c55d5579961fe64507f41b75f3c830 Mon Sep 17 00:00:00 2001 From: Alec Barber Date: Tue, 26 Nov 2024 05:14:36 +1100 Subject: [PATCH] celery: More precise typing for Celery.register_task(...) --- celery-stubs/app/base.pyi | 6 +----- tests/test_celery.py | 8 ++++++++ 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/celery-stubs/app/base.pyi b/celery-stubs/app/base.pyi index e198031..55bb431 100644 --- a/celery-stubs/app/base.pyi +++ b/celery-stubs/app/base.pyi @@ -268,11 +268,7 @@ class Celery: ) -> Callable[ [Callable[Concatenate[CeleryTask[_P, _R], _P], _R]], CeleryTask[_P, _R] ]: ... - def register_task( - self, - task: CeleryTask[Any, Any] | type[CeleryTask[Any, Any]], - **options: Any, - ) -> CeleryTask[Any, Any]: ... + def register_task(self, task: _T | type[_T], **options: Any) -> _T: ... def gen_task_name(self, name: str, module: str) -> str: ... def finalize(self, auto: bool = ...) -> None: ... def add_defaults(self, fun: Callable[[], dict[str, Any]]) -> None: ... diff --git a/tests/test_celery.py b/tests/test_celery.py index e7d0616..0ae02cb 100644 --- a/tests/test_celery.py +++ b/tests/test_celery.py @@ -137,6 +137,14 @@ def add_3(x: int, y: int) -> None: add_3.foo() +def test_celery_register_task() -> None: + bar = app_2.register_task(MyTask()) + bar.foo() + + baz = app_2.register_task(MyTask) + baz.foo() + + def test_celery_calling_task() -> None: signature("tasks.add", args=(2, 2), countdown=10)