Merge pull request #119 from mkurnikov/subclass-queryset-proper-typing

Allow to subclass queryset without loss of typing

README.md

@@ -42,6 +42,11 @@ django_settings_module = mysettings
 ```
 where `mysettings` is a value of `DJANGO_SETTINGS_MODULE` (with or without quotes)
 
+Do you have trouble with mypy / the django plugin not finding your settings module? Try adding the root path of your project to your PYTHONPATH environment variable. If you use pipenv you can add the following to an `.env` file in your project root which pipenv will run automatically before executing any commands.:
+```
+PYTHONPATH=${PYTHONPATH}:${PWD}
+```
+
 New implementation uses Django runtime to extract models information, so it will crash, if your installed apps `models.py` is not correct. For this same reason, you cannot use `reveal_type` inside global scope of any Python file that will be executed for `django.setup()`. 
 
 In other words, if your `manage.py runserver` crashes, mypy will crash too. 

dev-requirements.txt

@@ -1,5 +1,6 @@
 black
 pytest-mypy-plugins==1.0.3
+psycopg2
 flake8
 isort==4.3.4
 -e .

mypy_django_plugin/django/context.py

@@ -3,7 +3,6 @@ from collections import defaultdict
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Type
 
-from django.contrib.postgres.fields import ArrayField
 from django.core.exceptions import FieldError
 from django.db.models.base import Model
 from django.db.models.fields import AutoField, CharField, Field
@@ -18,6 +17,12 @@ from mypy.types import TypeOfAny
 
 from mypy_django_plugin.lib import helpers
 
+try:
+    from django.contrib.postgres.fields import ArrayField
+except ImportError:
+    class ArrayField:  # type: ignore
+        pass
+
 if TYPE_CHECKING:
     from django.apps.registry import Apps  # noqa: F401
     from django.conf import LazySettings  # noqa: F401

mypy_django_plugin/transformers/meta.py

@@ -5,7 +5,7 @@ from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny
 
 from mypy_django_plugin.django.context import DjangoContext
-from mypy_django_plugin.lib import fullnames, helpers
+from mypy_django_plugin.lib import helpers
 
 
 def _get_field_instance(ctx: MethodContext, field_fullname: str) -> MypyType:
@@ -20,21 +20,25 @@ def return_proper_field_type_from_get_field(ctx: MethodContext, django_context:
     # Options instance
     assert isinstance(ctx.type, Instance)
 
+    # bail if list of generic params is empty
+    if len(ctx.type.args) == 0:
+        return ctx.default_return_type
+
     model_type = ctx.type.args[0]
     if not isinstance(model_type, Instance):
-        return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
+        return ctx.default_return_type
 
     model_cls = django_context.get_model_class_by_fullname(model_type.type.fullname())
     if model_cls is None:
-        return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
+        return ctx.default_return_type
 
     field_name_expr = helpers.get_call_argument_by_name(ctx, 'field_name')
     if field_name_expr is None:
-        return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
+        return ctx.default_return_type
 
     field_name = helpers.resolve_string_attribute_value(field_name_expr, ctx, django_context)
     if field_name is None:
-        return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
+        return ctx.default_return_type
 
     try:
         field = model_cls._meta.get_field(field_name)

mypy_django_plugin/transformers/querysets.py

@@ -3,6 +3,7 @@ from typing import List, Optional, Sequence, Type, Union, cast
 
 from django.core.exceptions import FieldError
 from django.db.models.base import Model
+from django.db.models.fields.related import RelatedField
 from mypy.newsemanal.typeanal import TypeAnalyser
 from mypy.nodes import Expression, NameExpr, TypeInfo
 from mypy.plugin import AnalyzeTypeContext, FunctionContext, MethodContext
@@ -14,6 +15,15 @@ from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
 
 
+def _extract_model_type_from_queryset(queryset_type: Instance) -> Optional[Instance]:
+    for base_type in [queryset_type, *queryset_type.type.bases]:
+        if (len(base_type.args)
+                and isinstance(base_type.args[0], Instance)
+                and base_type.args[0].type.has_base(fullnames.MODEL_CLASS_FULLNAME)):
+            return base_type.args[0]
+    return None
+
+
 def determine_proper_manager_type(ctx: FunctionContext) -> MypyType:
     default_return_type = ctx.default_return_type
     assert isinstance(default_return_type, Instance)
@@ -34,6 +44,9 @@ def get_field_type_from_lookup(ctx: MethodContext, django_context: DjangoContext
         ctx.api.fail(exc.args[0], ctx.context)
         return None
 
+    if isinstance(lookup_field, RelatedField) and lookup_field.column == lookup:
+        lookup_field = django_context.get_primary_key_field(lookup_field.related_model)
+
     field_get_type = django_context.fields_context.get_field_get_type(helpers.get_typechecker_api(ctx),
                                                                       lookup_field, method=method)
     return field_get_type
@@ -94,11 +107,10 @@ def extract_proper_type_queryset_values_list(ctx: MethodContext, django_context:
     assert isinstance(ctx.type, Instance)
     assert isinstance(ctx.default_return_type, Instance)
 
-    # bail if queryset of Any or other non-instances
-    if not isinstance(ctx.type.args[0], Instance):
+    model_type = _extract_model_type_from_queryset(ctx.type)
+    if model_type is None:
         return AnyType(TypeOfAny.from_omitted_generics)
 
-    model_type = ctx.type.args[0]
     model_cls = django_context.get_model_class_by_fullname(model_type.type.fullname())
     if model_cls is None:
         return ctx.default_return_type
@@ -144,11 +156,10 @@ def extract_proper_type_queryset_values(ctx: MethodContext, django_context: Djan
     assert isinstance(ctx.type, Instance)
     assert isinstance(ctx.default_return_type, Instance)
 
-    # if queryset of non-instance type
-    if not isinstance(ctx.type.args[0], Instance):
+    model_type = _extract_model_type_from_queryset(ctx.type)
+    if model_type is None:
         return AnyType(TypeOfAny.from_omitted_generics)
 
-    model_type = ctx.type.args[0]
     model_cls = django_context.get_model_class_by_fullname(model_type.type.fullname())
     if model_cls is None:
         return ctx.default_return_type

setup.py

@@ -24,13 +24,11 @@ dependencies = [
     'mypy>=0.720,<0.730',
     'typing-extensions',
     'django',
-    # depends on psycopg2 because of Postgres' ArrayField support
-    'psycopg2'
 ]
 
 setup(
     name="django-stubs",
-    version="1.0.0",
+    version="1.0.2",
     description='Mypy stubs for Django',
     long_description=readme,
     long_description_content_type='text/markdown',

test-data/typecheck/managers/querysets/test_values_list.yml

@@ -188,3 +188,38 @@
                 class Entry(models.Model):
                     text = models.CharField(max_length=100)
                     blog = models.ForeignKey(Blog, on_delete=models.CASCADE)
+
+-   case: values_list_flat_true_with_ids
+    main: |
+        from myapp.models import Blog, Publisher
+        reveal_type(Blog.objects.values_list('id', flat=True))  # N: Revealed type is 'django.db.models.query.QuerySet[myapp.models.Blog, builtins.int]'
+        reveal_type(Blog.objects.values_list('publisher_id', flat=True))  # N: Revealed type is 'django.db.models.query.QuerySet[myapp.models.Blog, builtins.int]'
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class Publisher(models.Model):
+                    pass
+                class Blog(models.Model):
+                    publisher = models.ForeignKey(to=Publisher, on_delete=models.CASCADE)
+
+-   case: subclass_of_queryset_has_proper_typings_on_methods
+    main: |
+        from myapp.models import TransactionQuerySet
+        reveal_type(TransactionQuerySet())  # N: Revealed type is 'myapp.models.TransactionQuerySet'
+        reveal_type(TransactionQuerySet().values())  # N: Revealed type is 'django.db.models.query.QuerySet[myapp.models.Transaction, TypedDict({'id': builtins.int, 'total': builtins.int})]'
+        reveal_type(TransactionQuerySet().values_list())  # N: Revealed type is 'django.db.models.query.QuerySet[myapp.models.Transaction, Tuple[builtins.int, builtins.int]]'
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class TransactionQuerySet(models.QuerySet['Transaction']):
+                    pass
+                class Transaction(models.Model):
+                    total = models.IntegerField()