Don't modify return values of manager methods in place, modify a copy…

… instead (#2345)

mypy_django_plugin/transformers/managers.py

@@ -1,6 +1,7 @@
 from typing import Final, Optional, Union
 
 from mypy.checker import TypeChecker
+from mypy.copytype import copy_type
 from mypy.nodes import (
     GDEF,
     CallExpr,
@@ -244,15 +245,18 @@ def _replace_type_var(ret_type: MypyType, to_replace: str, replace_by: MypyType)
     """
     if isinstance(ret_type, TypeVarType) and ret_type.fullname == to_replace:
         return replace_by
-    elif isinstance((instance := get_proper_type(ret_type)), Instance):
+
+    ret_type = copy_type(get_proper_type(ret_type))
+
+    if isinstance(ret_type, Instance):
         # Since it is an instance, recursively find the type var for all its args.
-        instance.args = tuple(_replace_type_var(item, to_replace, replace_by) for item in instance.args)
-        return instance
+        ret_type.args = tuple(_replace_type_var(item, to_replace, replace_by) for item in ret_type.args)
+        return ret_type
 
-    if isinstance(ret_type, ProperType) and hasattr(ret_type, "item"):
+    if hasattr(ret_type, "item"):
         # For example TypeType has an item. find the type_var for this item
         ret_type.item = _replace_type_var(ret_type.item, to_replace, replace_by)
-    if isinstance(ret_type, ProperType) and hasattr(ret_type, "items"):
+    if hasattr(ret_type, "items"):
         # For example TypeList has items. find recursively type_var for its items
         ret_type.items = [_replace_type_var(item, to_replace, replace_by) for item in ret_type.items]
     return ret_type

tests/typecheck/managers/querysets/test_as_manager.yml

@@ -236,6 +236,7 @@
         reveal_type(MyModel.objects.dummy_override())  # N: Revealed type is "myapp.models.MyModel"
         reveal_type(MyModel.objects.example_mixin(MyModel()))  # N: Revealed type is "myapp.models.MyModel"
         reveal_type(MyModel.objects.example_other_mixin())  # N: Revealed type is "myapp.models.MyModel"
+        reveal_type(MyModel.objects.example_union())  # N: Revealed type is "Union[myapp.models.MyModel, builtins.list[myapp.models.MyModel]]"
         reveal_type(MyOtherModel.objects.example())  # N: Revealed type is "myapp.models.MyOtherModel"
         reveal_type(MyOtherModel.objects.example_2())  # N: Revealed type is "myapp.models.MyOtherModel"
         reveal_type(MyOtherModel.objects.override())  # N: Revealed type is "myapp.models.MyOtherModel"
@@ -245,13 +246,14 @@
         reveal_type(MyOtherModel.objects.example_other_mixin())  # N: Revealed type is "myapp.models.MyOtherModel"
         reveal_type(MyOtherModel.objects.test_self())  # N: Revealed type is "myapp.models._MyModelQuerySet2[myapp.models.MyOtherModel]"
         reveal_type(MyOtherModel.objects.test_sub_self())  # N: Revealed type is "myapp.models._MyModelQuerySet2[myapp.models.MyOtherModel]"
+        reveal_type(MyOtherModel.objects.example_union())  # N: Revealed type is "Union[myapp.models.MyOtherModel, builtins.list[myapp.models.MyOtherModel]]"
     installed_apps:
         - myapp
     files:
         -   path: myapp/__init__.py
         -   path: myapp/models.py
             content: |
-              from typing import TypeVar, Generic
+              from typing import TypeVar, Generic, Union, List
               from django.db import models
               from typing_extensions import Self
 
@@ -263,6 +265,7 @@
 
               class OtherMixin(models.QuerySet[T]):
                   def example_other_mixin(self) -> T: ...
+                  def example_union(self) -> Union[T, List[T]]: ...
 
               class _MyModelQuerySet(OtherMixin[T], models.QuerySet[T], Generic[T]):
                   def example(self) -> T: ...