Fix crash using .values_list("pk") on abstract model (#1355)

Author: flaeppe

mypy_django_plugin/django/context.py

@@ -2,11 +2,12 @@
 import sys
 from collections import defaultdict
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, Optional, Set, Tuple, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, Optional, Sequence, Set, Tuple, Type, Union
 
 from django.core.exceptions import FieldError
 from django.db import models
 from django.db.models.base import Model
+from django.db.models.expressions import Expression
 from django.db.models.fields import AutoField, CharField, Field
 from django.db.models.fields.related import ForeignKey, RelatedField
 from django.db.models.fields.reverse_related import ForeignObjectRel
@@ -19,6 +20,7 @@
 from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny, UnionType
+from typing_extensions import Literal
 
 from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.lib.fullnames import WITH_ANNOTATIONS_FULLNAME
@@ -375,30 +377,40 @@ def _resolve_field_from_parts(
         assert isinstance(field, (Field, ForeignObjectRel))
         return field
 
-    def resolve_lookup_into_field(
+    def solve_lookup_type(
         self, model_cls: Type[Model], lookup: str
-    ) -> Union["Field[Any, Any]", ForeignObjectRel]:
+    ) -> Optional[Tuple[Sequence[str], Sequence[str], Union[Expression, Literal[False]]]]:
         query = Query(model_cls)
-        lookup_parts, field_parts, is_expression = query.solve_lookup_type(lookup)
+        if (lookup == "pk" or lookup.startswith("pk__")) and query.get_meta().pk is None:
+            # Primary key lookup when no primary key field is found, model is presumably
+            # abstract and we can't say anything about 'pk'.
+            return None
+        return query.solve_lookup_type(lookup)
 
+    def resolve_lookup_into_field(
+        self, model_cls: Type[Model], lookup: str
+    ) -> Union["Field[Any, Any]", ForeignObjectRel, None]:
+        solved_lookup = self.solve_lookup_type(model_cls, lookup)
+        if solved_lookup is None:
+            return None
+        lookup_parts, field_parts, is_expression = solved_lookup
         if lookup_parts:
             raise LookupsAreUnsupported()
         return self._resolve_field_from_parts(field_parts, model_cls)
 
     def resolve_lookup_expected_type(self, ctx: MethodContext, model_cls: Type[Model], lookup: str) -> MypyType:
-        query = Query(model_cls)
-        if lookup == "pk" or lookup.startswith("pk__") and query.get_meta().pk is None:
-            # Primary key lookup when no primary key field is found, model is presumably
-            # abstract and we can't say anything about 'pk'.
-            return AnyType(TypeOfAny.implementation_artifact)
         try:
-            lookup_parts, field_parts, is_expression = query.solve_lookup_type(lookup)
-            if is_expression:
-                return AnyType(TypeOfAny.explicit)
+            solved_lookup = self.solve_lookup_type(model_cls, lookup)
         except FieldError as exc:
             ctx.api.fail(exc.args[0], ctx.context)
             return AnyType(TypeOfAny.from_error)
 
+        if solved_lookup is None:
+            return AnyType(TypeOfAny.implementation_artifact)
+        lookup_parts, field_parts, is_expression = solved_lookup
+        if is_expression:
+            return AnyType(TypeOfAny.explicit)
+
         field = self._resolve_field_from_parts(field_parts, model_cls)
 
         lookup_cls = None

mypy_django_plugin/transformers/querysets.py

@@ -58,7 +58,9 @@ def get_field_type_from_lookup(
     except LookupsAreUnsupported:
         return AnyType(TypeOfAny.explicit)
 
-    if (isinstance(lookup_field, RelatedField) and lookup_field.column == lookup) or isinstance(
+    if lookup_field is None:
+        return AnyType(TypeOfAny.implementation_artifact)
+    elif (isinstance(lookup_field, RelatedField) and lookup_field.column == lookup) or isinstance(
         lookup_field, ForeignObjectRel
     ):
         related_model_cls = django_context.get_field_related_model_cls(lookup_field)

tests/typecheck/models/test_abstract.yml

@@ -7,3 +7,29 @@
     AbstractUser.objects.get(pkey=1)  # ER: Cannot resolve keyword 'pkey' into field..*
   installed_apps:
     - django.contrib.auth
+
+- case: test_fetch_pk_with_custom_manager_on_abstract_model
+  main: |
+    from myapp.models import MyModel
+  installed_apps:
+    - myapp
+  files:
+    -   path: myapp/__init__.py
+    -   path: myapp/models.py
+        content: |
+            from django.db import models
+
+            class BaseManager(models.Manager):
+                pass
+
+            class BaseModel(models.Model):
+                objects = BaseManager()
+
+                class Meta:
+                    abstract = True
+
+                def lock(self) -> None:
+                    reveal_type(type(self).objects.values_list("pk"))  # N: Revealed type is "django.db.models.query._QuerySet[myapp.models.BaseModel, Tuple[Any]]"
+
+            class MyModel(BaseModel):
+                field = models.IntegerField()