diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index dfe02cf..aad8a42 100644 --- a/mypy_django_plugin/main.py +++ b/mypy_django_plugin/main.py @@ -64,13 +64,20 @@ def determine_proper_manager_type(ctx: FunctionContext) -> Type: if not isinstance(ret, Instance): return ret + has_manager_base = False for i, base in enumerate(ret.type.bases): if base.type.fullname() in {helpers.MANAGER_CLASS_FULLNAME, helpers.RELATED_MANAGER_CLASS_FULLNAME, helpers.BASE_MANAGER_CLASS_FULLNAME}: - ret.type.bases[i] = Instance(base.type, [Instance(outer_model_info, [])]) - return ret - return ret + has_manager_base = True + break + + if has_manager_base: + # Fill in the manager's type argument from the outer model + new_type_args = [Instance(outer_model_info, [])] + return helpers.reparametrize_instance(ret, new_type_args) + else: + return ret def return_user_model_hook(ctx: FunctionContext) -> Type: diff --git a/scripts/typecheck_tests.py b/scripts/typecheck_tests.py index 4e2bf3d..e59d185 100644 --- a/scripts/typecheck_tests.py +++ b/scripts/typecheck_tests.py @@ -199,9 +199,9 @@ IGNORED_ERRORS = { ], 'get_object_or_404': [ 'Argument 1 to "get_object_or_404" has incompatible type "str"; ' - + 'expected "Union[Type[], Manager[], QuerySet[]]"', + + 'expected "Union[Type[], QuerySet[]]"', 'Argument 1 to "get_list_or_404" has incompatible type "List[Type[Article]]"; ' - + 'expected "Union[Type[], Manager[], QuerySet[]]"', + + 'expected "Union[Type[], QuerySet[]]"', 'CustomClass' ], 'get_or_create': [ diff --git a/test-data/typecheck/managers.test b/test-data/typecheck/managers.test index 88ad066..527671f 100644 --- a/test-data/typecheck/managers.test +++ b/test-data/typecheck/managers.test @@ -136,4 +136,66 @@ class AbstractBase2(models.Model): class Child(AbstractBase1, AbstractBase2): pass +[out] + +[CASE managers_from_unrelated_models_dont_interfere] + +from django.db import models + +# Normal scenario where one model has a manager with an annotation of the same type as the model +class UnrelatedModel(models.Model): + objects = models.Manager[UnrelatedModel]() + +class MyModel(models.Model): + pass + +reveal_type(UnrelatedModel.objects) # E: Revealed type is 'django.db.models.manager.Manager[main.UnrelatedModel]' +reveal_type(UnrelatedModel.objects.first()) # E: Revealed type is 'Union[main.UnrelatedModel*, None]' + +reveal_type(MyModel.objects) # E: Revealed type is 'django.db.models.manager.Manager[main.MyModel]' +reveal_type(MyModel.objects.first()) # E: Revealed type is 'Union[main.MyModel*, None]' + +# Possible to specify objects without explicit annotation of models.Manager() +class UnrelatedModel2(models.Model): + objects = models.Manager() + +class MyModel2(models.Model): + pass + +reveal_type(UnrelatedModel2.objects) # E: Revealed type is 'django.db.models.manager.Manager[main.UnrelatedModel2]' +reveal_type(UnrelatedModel2.objects.first()) # E: Revealed type is 'Union[main.UnrelatedModel2*, None]' + +reveal_type(MyModel2.objects) # E: Revealed type is 'django.db.models.manager.Manager[main.MyModel2]' +reveal_type(MyModel2.objects.first()) # E: Revealed type is 'Union[main.MyModel2*, None]' + + +# Inheritance works +class ParentOfMyModel3(models.Model): + objects = models.Manager() + +class MyModel3(ParentOfMyModel3): + pass + + +reveal_type(ParentOfMyModel3.objects) # E: Revealed type is 'django.db.models.manager.Manager[main.ParentOfMyModel3]' +reveal_type(ParentOfMyModel3.objects.first()) # E: Revealed type is 'Union[main.ParentOfMyModel3*, None]' + +reveal_type(MyModel3.objects) # E: Revealed type is 'django.db.models.manager.Manager[main.MyModel3]' +reveal_type(MyModel3.objects.first()) # E: Revealed type is 'Union[main.MyModel3*, None]' + + +# Inheritance works with explicit objects in child +class ParentOfMyModel4(models.Model): + objects = models.Manager() + +class MyModel4(ParentOfMyModel4): + objects = models.Manager[MyModel4]() + +reveal_type(ParentOfMyModel4.objects) # E: Revealed type is 'django.db.models.manager.Manager[main.ParentOfMyModel4]' +reveal_type(ParentOfMyModel4.objects.first()) # E: Revealed type is 'Union[main.ParentOfMyModel4*, None]' + +reveal_type(MyModel4.objects) # E: Revealed type is 'django.db.models.manager.Manager[main.MyModel4]' +reveal_type(MyModel4.objects.first()) # E: Revealed type is 'Union[main.MyModel4*, None]' + + [out] \ No newline at end of file