From f182b39c91398019642159ba5f34e5c979c24356 Mon Sep 17 00:00:00 2001
From: Patrick Gingras <775.pg.12@gmail.com>
Date: Wed, 16 Jun 2021 03:15:32 -0400
Subject: [PATCH] Copy decorated queryset methods to manager too (#646)

* copy decorated queryset methods to manager too

* added test for from_manager with decorated queryset methods
---
 mypy_django_plugin/transformers/managers.py   | 13 +++++++----
 .../managers/querysets/test_from_queryset.yml | 22 +++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)

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()