QuerySet.annotate improvements (#398)

* QuerySet.annotate returns self-type. Attribute access falls back to Any.

- QuerySets that have an annotated model do not report errors during .filter() when called with invalid fields.
- QuerySets that have an annotated model return ordinary dict rather than TypedDict for .values()
- QuerySets that have an annotated model return Any rather than typed Tuple for .values_list()

* Fix .annotate so it reuses existing annotated types. Fixes error in typechecking Django testsuite.

* Fix self-typecheck error

* Fix flake8

* Fix case of .values/.values_list before .annotate.

* Extra ignores for Django 2.2 tests (false positives due to tests assuming QuerySet.first() won't return None)

Fix mypy self-check.

* More tests + more precise typing in case annotate called before values_list.

Cleanup tests.

* Test and fix annotate in combination with values/values_list with no params.

* Remove line that does nothing :)

* Formatting fixes

* Address code review

* Fix quoting in tests after mypy changed things

* Use Final

* Use typing_extensions.Final

* Fixes after ValuesQuerySet -> _ValuesQuerySet refactor. Still not passing tests yet.

* Fix inheritance of _ValuesQuerySet and remove unneeded type ignores.

This allows the test
"annotate_values_or_values_list_before_or_after_annotate_broadens_type"
to pass.

* Make it possible to annotate user code with "annotated models", using PEP 583 Annotated type.

* Add docs

* Make QuerySet[_T] an external alias to _QuerySet[_T, _T].

This currently has the drawback that error messages display the internal type _QuerySet, with both type arguments.

See also discussion on #661 and #608.

Fixes #635: QuerySet methods on Managers (like .all()) now return QuerySets rather than Managers.

Address code review by @sobolevn.

* Support passing TypedDicts to WithAnnotations

* Add an example of an error to README regarding WithAnnotations + TypedDict.

* Fix runtime behavior of ValuesQuerySet alias (you can't extend Any, for example).

Fix some edge case with from_queryset after QuerySet changed to be an
alias to _QuerySet. Can't make a minimal test case as this only occurred
on a large internal codebase.

* Fix issue when using from_queryset in some cases when having an argument with a type annotation on the QuerySet.

The mypy docstring on anal_type says not to call defer() after it.

mypy_django_plugin/transformers/models.py

@@ -1,19 +1,22 @@
-from typing import Dict, List, Optional, Type, cast
+from typing import Dict, List, Optional, Type, Union, cast
 
 from django.db.models.base import Model
 from django.db.models.fields import DateField, DateTimeField
 from django.db.models.fields.related import ForeignKey
 from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel, OneToOneRel
+from mypy.checker import TypeChecker
 from mypy.nodes import ARG_STAR2, Argument, Context, FuncDef, TypeInfo, Var
-from mypy.plugin import AttributeContext, ClassDefContext
+from mypy.plugin import AnalyzeTypeContext, AttributeContext, CheckerPluginInterface, ClassDefContext
 from mypy.plugins import common
 from mypy.semanal import SemanticAnalyzer
 from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
-from mypy.types import TypeOfAny
+from mypy.types import TypedDictType, TypeOfAny
 
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
+from mypy_django_plugin.lib.fullnames import ANNOTATIONS_FULLNAME, ANY_ATTR_ALLOWED_CLASS_FULLNAME
+from mypy_django_plugin.lib.helpers import add_new_class_for_module
 from mypy_django_plugin.transformers import fields
 from mypy_django_plugin.transformers.fields import get_field_descriptor_types
 
@@ -194,7 +197,6 @@ class AddManagers(ModelClassInitializer):
         for manager_name, manager in model_cls._meta.managers_map.items():
             manager_class_name = manager.__class__.__name__
             manager_fullname = helpers.get_class_fullname(manager.__class__)
-
             try:
                 manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname)
             except helpers.IncompleteDefnException as exc:
@@ -390,3 +392,76 @@ def set_auth_user_model_boolean_fields(ctx: AttributeContext, django_context: Dj
     boolinfo = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), bool)
     assert boolinfo is not None
     return Instance(boolinfo, [])
+
+
+def handle_annotated_type(ctx: AnalyzeTypeContext, django_context: DjangoContext) -> MypyType:
+    args = ctx.type.args
+    type_arg = ctx.api.analyze_type(args[0])
+    api = cast(SemanticAnalyzer, ctx.api.api)  # type: ignore
+
+    if not isinstance(type_arg, Instance):
+        return ctx.api.analyze_type(ctx.type)
+
+    fields_dict = None
+    if len(args) > 1:
+        second_arg_type = ctx.api.analyze_type(args[1])
+        if isinstance(second_arg_type, TypedDictType):
+            fields_dict = second_arg_type
+        elif isinstance(second_arg_type, Instance) and second_arg_type.type.fullname == ANNOTATIONS_FULLNAME:
+            annotations_type_arg = second_arg_type.args[0]
+            if isinstance(annotations_type_arg, TypedDictType):
+                fields_dict = annotations_type_arg
+            elif not isinstance(annotations_type_arg, AnyType):
+                ctx.api.fail("Only TypedDicts are supported as type arguments to Annotations", ctx.context)
+
+    return get_or_create_annotated_type(api, type_arg, fields_dict=fields_dict)
+
+
+def get_or_create_annotated_type(
+    api: Union[SemanticAnalyzer, CheckerPluginInterface], model_type: Instance, fields_dict: Optional[TypedDictType]
+) -> Instance:
+    """
+
+    Get or create the type for a model for which you getting/setting any attr is allowed.
+
+    The generated type is an subclass of the model and django._AnyAttrAllowed.
+    The generated type is placed in the django_stubs_ext module, with the name WithAnnotations[ModelName].
+    If the user wanted to annotate their code using this type, then this is the annotation they would use.
+    This is a bit of a hack to make a pretty type for error messages and which would make sense for users.
+    """
+    model_module_name = "django_stubs_ext"
+
+    if helpers.is_annotated_model_fullname(model_type.type.fullname):
+        # If it's already a generated class, we want to use the original model as a base
+        model_type = model_type.type.bases[0]
+
+    if fields_dict is not None:
+        type_name = f"WithAnnotations[{model_type.type.fullname}, {fields_dict}]"
+    else:
+        type_name = f"WithAnnotations[{model_type.type.fullname}]"
+
+    annotated_typeinfo = helpers.lookup_fully_qualified_typeinfo(
+        cast(TypeChecker, api), model_module_name + "." + type_name
+    )
+    if annotated_typeinfo is None:
+        model_module_file = api.modules[model_module_name]  # type: ignore
+
+        if isinstance(api, SemanticAnalyzer):
+            annotated_model_type = api.named_type_or_none(ANY_ATTR_ALLOWED_CLASS_FULLNAME, [])
+            assert annotated_model_type is not None
+        else:
+            annotated_model_type = api.named_generic_type(ANY_ATTR_ALLOWED_CLASS_FULLNAME, [])
+
+        annotated_typeinfo = add_new_class_for_module(
+            model_module_file,
+            type_name,
+            bases=[model_type] if fields_dict is not None else [model_type, annotated_model_type],
+            fields=fields_dict.items if fields_dict is not None else None,
+        )
+        if fields_dict is not None:
+            # To allow structural subtyping, make it a Protocol
+            annotated_typeinfo.is_protocol = True
+            # Save for later to easily find which field types were annotated
+            annotated_typeinfo.metadata["annotated_field_types"] = fields_dict.items
+    annotated_type = Instance(annotated_typeinfo, [])
+    return annotated_type