Mypy Celery plugin to patch the return type of shared_task and Celery.task decorators with a Django-specific task class.

mypy_celery_plugin.py

from collections.abc import Callable

from mypy.nodes import MypyFile
from mypy.plugin import FunctionContext, MethodContext, Plugin
from mypy.types import CallableType, Instance, Type, get_proper_type

DJANGO_TASK_FULLNAME = "celery.contrib.django.task.DjangoTask"


class CeleryPlugin(Plugin):
    def get_function_hook(
        self, fullname: str
    ) -> Callable[[FunctionContext], Type] | None:
        if fullname == "celery.app.shared_task":
            return self.celery_task_hook
        return None

    def get_method_hook(self, fullname: str) -> Callable[[MethodContext], Type] | None:
        if fullname == "celery.app.base.Celery.task":
            return self.celery_task_hook
        return None

    def celery_task_hook(self, context: FunctionContext | MethodContext) -> Type:
        sym = self.lookup_fully_qualified(DJANGO_TASK_FULLNAME)

        return_type = get_proper_type(context.default_return_type)
        if isinstance(return_type, CallableType):
            # decorator called with arguments
            return return_type.copy_modified(
                ret_type=Instance(sym.node, return_type.ret_type.args)
            )

        return Instance(sym.node, return_type.args)

    def get_additional_deps(self, file: MypyFile) -> list[tuple[int, str, int]]:
        if "celery" in file.fullname:
            return [(10, DJANGO_TASK_FULLNAME.rsplit(".", 1)[0], -1)]
        return []


def plugin(version: str) -> type[CeleryPlugin]:
    return CeleryPlugin