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

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

* linting

mypy_django_plugin/transformers/init_create.py

@@ -6,7 +6,7 @@ from mypy.types import Instance
 from mypy.types import Type as MypyType
 
 from mypy_django_plugin.django.context import DjangoContext
-from mypy_django_plugin.lib import helpers
+from mypy_django_plugin.lib import fullnames, helpers
 
 
 def get_actual_types(ctx: Union[MethodContext, FunctionContext],
@@ -30,6 +30,12 @@ def get_actual_types(ctx: Union[MethodContext, FunctionContext],
     return actual_types
 
 
+def check_types_compatible(ctx, *, expected_type, actual_type, error_message):
+    ctx.api.check_subtype(actual_type, expected_type,
+                          ctx.context, error_message,
+                          'got', 'expected')
+
+
 def typecheck_model_method(ctx: Union[FunctionContext, MethodContext], django_context: DjangoContext,
                            model_cls: Type[Model], method: str) -> MypyType:
     typechecker_api = helpers.get_typechecker_api(ctx)
@@ -42,11 +48,11 @@ def typecheck_model_method(ctx: Union[FunctionContext, MethodContext], django_co
                                                                            model_cls.__name__),
                          ctx.context)
             continue
-        typechecker_api.check_subtype(actual_type, expected_types[actual_name],
-                                      ctx.context,
-                                      'Incompatible type for "{}" of "{}"'.format(actual_name,
-                                                                                  model_cls.__name__),
-                                      'got', 'expected')
+        check_types_compatible(ctx,
+                               expected_type=expected_types[actual_name],
+                               actual_type=actual_type,
+                               error_message='Incompatible type for "{}" of "{}"'.format(actual_name,
+                                                                                         model_cls.__name__))
 
     return ctx.default_return_type
 
@@ -73,3 +79,40 @@ def redefine_and_typecheck_model_create(ctx: MethodContext, django_context: Djan
         return ctx.default_return_type
 
     return typecheck_model_method(ctx, django_context, model_cls, 'create')
+
+
+def typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
+    lookup_kwargs = ctx.arg_names[1]
+    provided_lookup_types = ctx.arg_types[1]
+
+    assert isinstance(ctx.type, Instance)
+
+    if not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
+        return ctx.default_return_type
+
+    model_cls_fullname = ctx.type.args[0].type.fullname()
+    model_cls = django_context.get_model_class_by_fullname(model_cls_fullname)
+    if model_cls is None:
+        return ctx.default_return_type
+
+    for lookup_kwarg, provided_type in zip(lookup_kwargs, provided_lookup_types):
+        if lookup_kwarg is None:
+            continue
+        # Combinables are not supported yet
+        if (isinstance(provided_type, Instance)
+                and provided_type.type.has_base('django.db.models.expressions.Combinable')):
+            continue
+
+        lookup_type = django_context.lookups_context.resolve_lookup_expected_type(ctx, model_cls, lookup_kwarg)
+        # Managers as provided_type is not supported yet
+        if (isinstance(provided_type, Instance)
+                and helpers.has_any_of_bases(provided_type.type, (fullnames.MANAGER_CLASS_FULLNAME,
+                                                                  fullnames.QUERYSET_CLASS_FULLNAME))):
+            return ctx.default_return_type
+
+        check_types_compatible(ctx,
+                               expected_type=lookup_type,
+                               actual_type=provided_type,
+                               error_message=f'Incompatible type for lookup {lookup_kwarg!r}:')
+
+    return ctx.default_return_type

mypy_django_plugin/transformers/querysets.py

@@ -10,7 +10,9 @@ from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny
 
-from mypy_django_plugin.django.context import DjangoContext
+from mypy_django_plugin.django.context import (
+    DjangoContext, LookupsAreUnsupported,
+)
 from mypy_django_plugin.lib import fullnames, helpers
 
 
@@ -38,10 +40,12 @@ def determine_proper_manager_type(ctx: FunctionContext) -> MypyType:
 def get_field_type_from_lookup(ctx: MethodContext, django_context: DjangoContext, model_cls: Type[Model],
                                *, method: str, lookup: str) -> Optional[MypyType]:
     try:
-        lookup_field = django_context.lookups_context.resolve_lookup(model_cls, lookup)
+        lookup_field = django_context.lookups_context.resolve_lookup_info_field(model_cls, lookup)
     except FieldError as exc:
         ctx.api.fail(exc.args[0], ctx.context)
         return None
+    except LookupsAreUnsupported:
+        return AnyType(TypeOfAny.explicit)
 
     if isinstance(lookup_field, RelatedField) and lookup_field.column == lookup:
         related_model_cls = django_context.fields_context.get_related_model_cls(lookup_field)