add support for typechecking of filter/get/exclude arguments (#183)

* add support for typechecking of filter/get/exclude arguments

* linting

Author: mkurnikov

django-stubs/db/models/fields/__init__.pyi

@@ -43,6 +43,7 @@ _GT = TypeVar("_GT")
 class Field(RegisterLookupMixin, Generic[_ST, _GT]):
     _pyi_private_set_type: Any
     _pyi_private_get_type: Any
+    _pyi_lookup_exact_type: Any
 
     widget: Widget
     help_text: str
@@ -131,6 +132,7 @@ class Field(RegisterLookupMixin, Generic[_ST, _GT]):
 class IntegerField(Field[_ST, _GT]):
     _pyi_private_set_type: Union[float, int, str, Combinable]
     _pyi_private_get_type: int
+    _pyi_lookup_exact_type: int
 
 class PositiveIntegerRelDbTypeMixin:
     def rel_db_type(self, connection: Any): ...
@@ -143,10 +145,12 @@ class BigIntegerField(IntegerField[_ST, _GT]): ...
 class FloatField(Field[_ST, _GT]):
     _pyi_private_set_type: Union[float, int, str, Combinable]
     _pyi_private_get_type: float
+    _pyi_lookup_exact_type: float
 
 class DecimalField(Field[_ST, _GT]):
     _pyi_private_set_type: Union[str, float, decimal.Decimal, Combinable]
     _pyi_private_get_type: decimal.Decimal
+    _pyi_lookup_exact_type: Union[str, decimal.Decimal]
     # attributes
     max_digits: int = ...
     decimal_places: int = ...
@@ -176,10 +180,13 @@ class DecimalField(Field[_ST, _GT]):
 class AutoField(Field[_ST, _GT]):
     _pyi_private_set_type: Union[Combinable, int, str]
     _pyi_private_get_type: int
+    _pyi_lookup_exact_type: int
 
 class CharField(Field[_ST, _GT]):
     _pyi_private_set_type: Union[str, int, Combinable]
     _pyi_private_get_type: str
+    # objects are converted to string before comparison
+    _pyi_lookup_exact_type: Any
     def __init__(
         self,
         verbose_name: Optional[Union[str, bytes]] = ...,
@@ -238,14 +245,18 @@ class URLField(CharField[_ST, _GT]): ...
 class TextField(Field[_ST, _GT]):
     _pyi_private_set_type: Union[str, Combinable]
     _pyi_private_get_type: str
+    # objects are converted to string before comparison
+    _pyi_lookup_exact_type: Any
 
 class BooleanField(Field[_ST, _GT]):
     _pyi_private_set_type: Union[bool, Combinable]
     _pyi_private_get_type: bool
+    _pyi_lookup_exact_type: bool
 
 class NullBooleanField(Field[_ST, _GT]):
     _pyi_private_set_type: Optional[Union[bool, Combinable]]
     _pyi_private_get_type: Optional[bool]
+    _pyi_lookup_exact_type: Optional[bool]
 
 class IPAddressField(Field[_ST, _GT]):
     _pyi_private_set_type: Union[str, Combinable]
@@ -286,6 +297,7 @@ class DateTimeCheckMixin: ...
 class DateField(DateTimeCheckMixin, Field[_ST, _GT]):
     _pyi_private_set_type: Union[str, date, Combinable]
     _pyi_private_get_type: date
+    _pyi_lookup_exact_type: Union[str, date]
     def __init__(
         self,
         verbose_name: Optional[Union[str, bytes]] = ...,
@@ -338,6 +350,7 @@ class TimeField(DateTimeCheckMixin, Field[_ST, _GT]):
 
 class DateTimeField(DateField[_ST, _GT]):
     _pyi_private_get_type: datetime
+    _pyi_lookup_exact_type: Union[str, datetime]
 
 class UUIDField(Field[_ST, _GT]):
     _pyi_private_set_type: Union[str, uuid.UUID]

django-stubs/db/models/lookups.pyi

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Any, Iterable, List, Mapping, Optional, Tuple, Type, Union
+from typing import Any, Iterable, List, Optional, Tuple, Type, Union, Mapping, TypeVar, Generic
 
 from django.db.backends.sqlite3.base import DatabaseWrapper
 from django.db.models.expressions import Expression, Func
@@ -10,7 +10,9 @@ from django.utils.safestring import SafeText
 
 from django.db.models.fields import TextField, related_lookups
 
-class Lookup:
+_T = TypeVar("_T")
+
+class Lookup(Generic[_T]):
     lookup_name: str = ...
     prepare_rhs: bool = ...
     can_use_none_as_rhs: bool = ...
@@ -47,7 +49,7 @@ class Transform(RegisterLookupMixin, Func):
     def lhs(self) -> Expression: ...
     def get_bilateral_transforms(self) -> List[Type[Transform]]: ...
 
-class BuiltinLookup(Lookup):
+class BuiltinLookup(Lookup[_T]):
     def get_rhs_op(self, connection: DatabaseWrapper, rhs: str) -> str: ...
 
 class FieldGetDbPrepValueMixin:
@@ -62,21 +64,21 @@ class FieldGetDbPrepValueIterableMixin(FieldGetDbPrepValueMixin):
 class Exact(FieldGetDbPrepValueMixin, BuiltinLookup): ...
 class IExact(BuiltinLookup): ...
 class GreaterThan(FieldGetDbPrepValueMixin, BuiltinLookup): ...
-class GreaterThanOrEqual(FieldGetDbPrepValueMixin, BuiltinLookup): ...
-class LessThan(FieldGetDbPrepValueMixin, BuiltinLookup): ...
+class GreaterThanOrEqual(FieldGetDbPrepValueMixin, BuiltinLookup[_T]): ...
+class LessThan(FieldGetDbPrepValueMixin, BuiltinLookup[_T]): ...
 class LessThanOrEqual(FieldGetDbPrepValueMixin, BuiltinLookup): ...
 
 class IntegerFieldFloatRounding:
     rhs: Any = ...
     def get_prep_lookup(self) -> Any: ...
 
-class IntegerGreaterThanOrEqual(IntegerFieldFloatRounding, GreaterThanOrEqual): ...
-class IntegerLessThan(IntegerFieldFloatRounding, LessThan): ...
+class IntegerGreaterThanOrEqual(IntegerFieldFloatRounding, GreaterThanOrEqual[Union[int, float]]): ...
+class IntegerLessThan(IntegerFieldFloatRounding, LessThan[Union[int, float]]): ...
 
 class In(FieldGetDbPrepValueIterableMixin, BuiltinLookup):
     def split_parameter_list_as_sql(self, compiler: Any, connection: Any): ...
 
-class PatternLookup(BuiltinLookup):
+class PatternLookup(BuiltinLookup[str]):
     param_pattern: str = ...
 
 class Contains(PatternLookup): ...
@@ -86,8 +88,8 @@ class IStartsWith(StartsWith): ...
 class EndsWith(PatternLookup): ...
 class IEndsWith(EndsWith): ...
 class Range(FieldGetDbPrepValueIterableMixin, BuiltinLookup): ...
-class IsNull(BuiltinLookup): ...
-class Regex(BuiltinLookup): ...
+class IsNull(BuiltinLookup[bool]): ...
+class Regex(BuiltinLookup[str]): ...
 class IRegex(Regex): ...
 
 class YearLookup(Lookup):

mypy_django_plugin/django/context.py

@@ -2,7 +2,7 @@
 from collections import defaultdict
 from contextlib import contextmanager
 from typing import (
-    TYPE_CHECKING, Dict, Iterator, Optional, Set, Tuple, Type, Union,
+    TYPE_CHECKING, Dict, Iterable, Iterator, Optional, Set, Tuple, Type, Union,
 )
 
 from django.core.exceptions import FieldError
@@ -11,14 +11,16 @@
 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
+from django.db.models.lookups import Exact
 from django.db.models.sql.query import Query
 from django.utils.functional import cached_property
 from mypy.checker import TypeChecker
+from mypy.plugin import MethodContext
 from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
-from mypy.types import TypeOfAny
+from mypy.types import TypeOfAny, UnionType
 
-from mypy_django_plugin.lib import helpers
+from mypy_django_plugin.lib import fullnames, helpers
 
 try:
     from django.contrib.postgres.fields import ArrayField
@@ -153,33 +155,87 @@ def get_related_model_cls(self, field: Union[RelatedField, ForeignObjectRel]) ->
         return related_model_cls
 
 
+class LookupsAreUnsupported(Exception):
+    pass
+
+
 class DjangoLookupsContext:
     def __init__(self, django_context: 'DjangoContext'):
         self.django_context = django_context
 
-    def resolve_lookup(self, model_cls: Type[Model], lookup: str) -> Field:
+    def _resolve_field_from_parts(self, field_parts: Iterable[str], model_cls: Type[Model]) -> Field:
+        currently_observed_model = model_cls
+        field = None
+        for field_part in field_parts:
+            if field_part == 'pk':
+                field = self.django_context.get_primary_key_field(currently_observed_model)
+                continue
+
+            field = currently_observed_model._meta.get_field(field_part)
+            if isinstance(field, RelatedField):
+                currently_observed_model = field.related_model
+                model_name = currently_observed_model._meta.model_name
+                if (model_name is not None
+                        and field_part == (model_name + '_id')):
+                    field = self.django_context.get_primary_key_field(currently_observed_model)
+
+            if isinstance(field, ForeignObjectRel):
+                currently_observed_model = field.related_model
+
+        assert field is not None
+        return field
+
+    def resolve_lookup_info_field(self, model_cls: Type[Model], lookup: str) -> Field:
         query = Query(model_cls)
         lookup_parts, field_parts, is_expression = query.solve_lookup_type(lookup)
         if lookup_parts:
-            raise FieldError('Lookups not supported yet')
+            raise LookupsAreUnsupported()
 
-        currently_observed_model = model_cls
-        current_field = None
-        for field_part in field_parts:
-            if field_part == 'pk':
-                return self.django_context.get_primary_key_field(currently_observed_model)
+        return self._resolve_field_from_parts(field_parts, model_cls)
 
-            current_field = currently_observed_model._meta.get_field(field_part)
-            if not isinstance(current_field, (ForeignObjectRel, RelatedField)):
-                continue
+    def resolve_lookup_expected_type(self, ctx: MethodContext, model_cls: Type[Model], lookup: str) -> MypyType:
+        query = Query(model_cls)
+        try:
+            lookup_parts, field_parts, is_expression = query.solve_lookup_type(lookup)
+            if is_expression:
+                return AnyType(TypeOfAny.explicit)
+        except FieldError as exc:
+            ctx.api.fail(exc.args[0], ctx.context)
+            return AnyType(TypeOfAny.from_error)
+
+        field = self._resolve_field_from_parts(field_parts, model_cls)
+
+        lookup_cls = None
+        if lookup_parts:
+            lookup = lookup_parts[-1]
+            lookup_cls = field.get_lookup(lookup)
+            if lookup_cls is None:
+                # unknown lookup
+                return AnyType(TypeOfAny.explicit)
+
+        if lookup_cls is None or isinstance(lookup_cls, Exact):
+            return self.django_context.get_field_lookup_exact_type(helpers.get_typechecker_api(ctx), field)
+
+        assert lookup_cls is not None
 
-            currently_observed_model = self.django_context.fields_context.get_related_model_cls(current_field)
-            if isinstance(current_field, ForeignObjectRel):
-                current_field = self.django_context.get_primary_key_field(currently_observed_model)
+        lookup_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), lookup_cls)
+        if lookup_info is None:
+            return AnyType(TypeOfAny.explicit)
 
-        # if it is None, solve_lookup_type() will fail earlier
-        assert current_field is not None
-        return current_field
+        for lookup_base in helpers.iter_bases(lookup_info):
+            if lookup_base.args and isinstance(lookup_base.args[0], Instance):
+                lookup_type: MypyType = lookup_base.args[0]
+                # if it's Field, consider lookup_type a __get__ of current field
+                if (isinstance(lookup_type, Instance)
+                        and lookup_type.type.fullname() == fullnames.FIELD_FULLNAME):
+                    field_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), field.__class__)
+                    if field_info is None:
+                        return AnyType(TypeOfAny.explicit)
+                    lookup_type = helpers.get_private_descriptor_type(field_info, '_pyi_private_get_type',
+                                                                      is_nullable=field.null)
+                return lookup_type
+
+        return AnyType(TypeOfAny.explicit)
 
 
 class DjangoContext:
@@ -228,6 +284,27 @@ def get_model_relations(self, model_cls: Type[Model]) -> Iterator[ForeignObjectR
             if isinstance(field, ForeignObjectRel):
                 yield field
 
+    def get_field_lookup_exact_type(self, api: TypeChecker, field: Field) -> MypyType:
+        if isinstance(field, (RelatedField, ForeignObjectRel)):
+            related_model_cls = field.related_model
+            primary_key_field = self.get_primary_key_field(related_model_cls)
+            primary_key_type = self.fields_context.get_field_get_type(api, primary_key_field, method='init')
+
+            rel_model_info = helpers.lookup_class_typeinfo(api, related_model_cls)
+            if rel_model_info is None:
+                return AnyType(TypeOfAny.explicit)
+
+            model_and_primary_key_type = UnionType.make_union([Instance(rel_model_info, []),
+                                                               primary_key_type])
+            return helpers.make_optional(model_and_primary_key_type)
+            # return helpers.make_optional(Instance(rel_model_info, []))
+
+        field_info = helpers.lookup_class_typeinfo(api, field.__class__)
+        if field_info is None:
+            return AnyType(TypeOfAny.explicit)
+        return helpers.get_private_descriptor_type(field_info, '_pyi_lookup_exact_type',
+                                                   is_nullable=field.null)
+
     def get_primary_key_field(self, model_cls: Type[Model]) -> Field:
         for field in model_cls._meta.get_fields():
             if isinstance(field, Field):

mypy_django_plugin/lib/helpers.py

@@ -1,6 +1,11 @@
 from collections import OrderedDict
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union, cast
+from typing import (
+    TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Set, Union, cast,
+)
 
+from django.db.models.fields import Field
+from django.db.models.fields.related import RelatedField
+from django.db.models.fields.reverse_related import ForeignObjectRel
 from mypy import checker
 from mypy.checker import TypeChecker
 from mypy.mro import calculate_mro
@@ -115,29 +120,50 @@ def parse_bool(expr: Expression) -> Optional[bool]:
     return None
 
 
-def has_any_of_bases(info: TypeInfo, bases: Set[str]) -> bool:
+def has_any_of_bases(info: TypeInfo, bases: Iterable[str]) -> bool:
     for base_fullname in bases:
         if info.has_base(base_fullname):
             return True
     return False
 
 
+def iter_bases(info: TypeInfo) -> Iterator[Instance]:
+    for base in info.bases:
+        yield base
+        yield from iter_bases(base.type)
+
+
 def get_private_descriptor_type(type_info: TypeInfo, private_field_name: str, is_nullable: bool) -> MypyType:
     """ Return declared type of type_info's private_field_name (used for private Field attributes)"""
     sym = type_info.get(private_field_name)
     if sym is None:
-        return AnyType(TypeOfAny.unannotated)
+        return AnyType(TypeOfAny.explicit)
 
     node = sym.node
     if isinstance(node, Var):
         descriptor_type = node.type
         if descriptor_type is None:
-            return AnyType(TypeOfAny.unannotated)
+            return AnyType(TypeOfAny.explicit)
 
         if is_nullable:
             descriptor_type = make_optional(descriptor_type)
         return descriptor_type
-    return AnyType(TypeOfAny.unannotated)
+    return AnyType(TypeOfAny.explicit)
+
+
+def get_field_lookup_exact_type(api: TypeChecker, field: Field) -> MypyType:
+    if isinstance(field, (RelatedField, ForeignObjectRel)):
+        lookup_type_class = field.related_model
+        rel_model_info = lookup_class_typeinfo(api, lookup_type_class)
+        if rel_model_info is None:
+            return AnyType(TypeOfAny.from_error)
+        return make_optional(Instance(rel_model_info, []))
+
+    field_info = lookup_class_typeinfo(api, field.__class__)
+    if field_info is None:
+        return AnyType(TypeOfAny.explicit)
+    return get_private_descriptor_type(field_info, '_pyi_lookup_exact_type',
+                                       is_nullable=field.null)
 
 
 def get_nested_meta_node_for_current_class(info: TypeInfo) -> Optional[TypeInfo]:

mypy_django_plugin/main.py

@@ -209,6 +209,8 @@ def get_method_hook(self, fullname: str
         manager_classes = self._get_current_manager_bases()
         if class_fullname in manager_classes and method_name == 'create':
             return partial(init_create.redefine_and_typecheck_model_create, django_context=self.django_context)
+        if class_fullname in manager_classes and method_name in {'filter', 'get', 'exclude'}:
+            return partial(init_create.typecheck_queryset_filter, django_context=self.django_context)
         return None
 
     def get_base_class_hook(self, fullname: str