celery: More precise typing for Celery.register_task(...) (#168)

Resolves  #167 

`Celery.register_task` returns the task instance that was registered, which might be a subclass of `Task` [1]. We can accurately capture this behaviour using a `TypeVar`.

[1] https://github.com/celery/celery/blob/3630e467361009a8b3f3050807ed16503d5c4441/celery/app/base.py#L609

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: ...

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)