diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 976e4cd..608477b 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -1,4 +1,4 @@ -from mypy.nodes import GDEF, FuncDef, MemberExpr, NameExpr, RefExpr, StrExpr, SymbolTableNode, TypeInfo +from mypy.nodes import GDEF, Decorator, FuncDef, MemberExpr, NameExpr, RefExpr, StrExpr, SymbolTableNode, TypeInfo from mypy.plugin import ClassDefContext, DynamicClassDefContext from mypy.types import AnyType, Instance, TypeOfAny @@ -67,6 +67,11 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte break for name, sym in class_mro_info.names.items(): if isinstance(sym.node, FuncDef): - helpers.copy_method_to_another_class( - class_def_context, self_type, new_method_name=name, method_node=sym.node - ) + func_node = sym.node + elif isinstance(sym.node, Decorator): + func_node = sym.node.func + else: + continue + helpers.copy_method_to_another_class( + class_def_context, self_type, new_method_name=name, method_node=func_node + ) diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index e291ba8..53dceac 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -179,3 +179,25 @@ class BaseQuerySet(models.QuerySet): def base_queryset_method(self, param: Union[int, str]) -> NoReturn: raise ValueError + +- case: from_queryset_with_decorated_queryset_methods + main: | + from myapp.models import MyModel + reveal_type(MyModel().objects) # N: Revealed type is "myapp.models.MyModel_NewManager[myapp.models.MyModel]" + reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is "builtins.str" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models, transaction + + class ModelQuerySet(models.QuerySet): + @transaction.atomic + def queryset_method(self) -> str: + return 'hello' + + NewManager = models.Manager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + objects = NewManager()