diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py index 1ae6f3ca3..acd7fef67 100644 --- a/mypy_django_plugin/django/context.py +++ b/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 diff --git a/mypy_django_plugin/transformers/querysets.py b/mypy_django_plugin/transformers/querysets.py index 566101908..0dcabfdfc 100644 --- a/mypy_django_plugin/transformers/querysets.py +++ b/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) diff --git a/tests/typecheck/models/test_abstract.yml b/tests/typecheck/models/test_abstract.yml index e7303b5f1..19f3bfb51 100644 --- a/tests/typecheck/models/test_abstract.yml +++ b/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()