Fix mypy cache of WithAnnotation types

Author: kalekseev

mypy_django_plugin/transformers/models.py

@@ -1,3 +1,5 @@
+import codecs
+import pickle
 from typing import Dict, List, Optional, Type, Union, cast
 
 from django.db.models.base import Model
@@ -436,9 +438,9 @@ def get_or_create_annotated_type(
         model_type = model_type.type.bases[0]
 
     if fields_dict is not None:
-        type_name = f"WithAnnotations[{model_type.type.fullname}, {fields_dict}]"
+        type_name = f"WithAnnotations[{model_type.type.fullname.replace('.', '__')}, {fields_dict}]"
     else:
-        type_name = f"WithAnnotations[{model_type.type.fullname}]"
+        type_name = f"WithAnnotations[{model_type.type.fullname.replace('.', '__')}]"
 
     annotated_typeinfo = helpers.lookup_fully_qualified_typeinfo(
         cast(TypeChecker, api), model_module_name + "." + type_name
@@ -462,6 +464,19 @@ def get_or_create_annotated_type(
             # 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
+            store_annotate_field_types(annotated_typeinfo, fields_dict)
     annotated_type = Instance(annotated_typeinfo, [])
     return annotated_type
+
+
+def store_annotate_field_types(annotated_typeinfo, fields_dict: TypedDictType) -> None:
+    annotated_typeinfo.metadata["annotated_field_types"] = codecs.encode(
+        pickle.dumps(fields_dict.items), "base64"
+    ).decode("utf-8")
+
+
+def retrieve_annotated_field_types(annotated_typeinfo) -> Optional[Dict[str, AnyType]]:
+    field_dict = annotated_typeinfo.metadata.get("annotated_field_types")
+    if field_dict is None:
+        return field_dict
+    return pickle.loads(codecs.decode(field_dict.encode("utf-8"), "base64"))

mypy_django_plugin/transformers/querysets.py

@@ -15,7 +15,7 @@
 from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.lib.fullnames import ANY_ATTR_ALLOWED_CLASS_FULLNAME
 from mypy_django_plugin.lib.helpers import is_annotated_model_fullname
-from mypy_django_plugin.transformers.models import get_or_create_annotated_type
+from mypy_django_plugin.transformers.models import get_or_create_annotated_type, retrieve_annotated_field_types
 
 
 def _extract_model_type_from_queryset(queryset_type: Instance) -> Optional[Instance]:
@@ -212,7 +212,7 @@ def extract_proper_type_queryset_annotate(ctx: MethodContext, django_context: Dj
 
     api = ctx.api
 
-    field_types = model_type.type.metadata.get("annotated_field_types")
+    field_types = retrieve_annotated_field_types(model_type.type)
     kwargs = gather_kwargs(ctx)
     if kwargs:
         # For now, we don't try to resolve the output_field of the field would be, but use Any.

tests/typecheck/managers/querysets/test_annotate.yml

@@ -8,20 +8,20 @@
 
         unannotated_user = User.objects.get(id=1)
 
-        print(annotated_user.asdf) # E: "WithAnnotations[myapp.models.User, TypedDict({'foo': Any})]" has no attribute "asdf"
+        print(annotated_user.asdf) # E: "WithAnnotations[myapp__models__User, TypedDict({'foo': Any})]" has no attribute "asdf"
         print(unannotated_user.asdf) # E: "User" has no attribute "asdf"
 
         def func(user: Annotated[User, Annotations]) -> str:
             return user.asdf
 
-        func(unannotated_user) # E: Argument 1 to "func" has incompatible type "User"; expected "WithAnnotations[myapp.models.User]"
-        func(annotated_user) # E: Argument 1 to "func" has incompatible type "WithAnnotations[myapp.models.User, TypedDict({'foo': Any})]"; expected "WithAnnotations[myapp.models.User]"
+        func(unannotated_user) # E: Argument 1 to "func" has incompatible type "User"; expected "WithAnnotations[myapp__models__User]"
+        func(annotated_user) # E: Argument 1 to "func" has incompatible type "WithAnnotations[myapp__models__User, TypedDict({'foo': Any})]"; expected "WithAnnotations[myapp__models__User]"
 
         def func2(user: WithAnnotations[User]) -> str:
             return user.asdf
 
-        func2(unannotated_user) # E: Argument 1 to "func2" has incompatible type "User"; expected "WithAnnotations[myapp.models.User]"
-        func2(annotated_user) # E: Argument 1 to "func2" has incompatible type "WithAnnotations[myapp.models.User, TypedDict({'foo': Any})]"; expected "WithAnnotations[myapp.models.User]"
+        func2(unannotated_user) # E: Argument 1 to "func2" has incompatible type "User"; expected "WithAnnotations[myapp__models__User]"
+        func2(annotated_user) # E: Argument 1 to "func2" has incompatible type "WithAnnotations[myapp__models__User, TypedDict({'foo': Any})]"; expected "WithAnnotations[myapp__models__User]"
     installed_apps:
         - myapp
     files:
@@ -44,26 +44,26 @@
             foo: str
 
         def func(user: Annotated[User, Annotations[MyDict]]) -> str:
-            print(user.asdf) # E: "WithAnnotations[myapp.models.User, TypedDict('main.MyDict', {'foo': builtins.str})]" has no attribute "asdf"
+            print(user.asdf) # E: "WithAnnotations[myapp__models__User, TypedDict('main.MyDict', {'foo': builtins.str})]" has no attribute "asdf"
             return user.foo
 
         unannotated_user = User.objects.get(id=1)
         annotated_user = User.objects.annotate(foo=Value("")).get()
         other_annotated_user = User.objects.annotate(other=Value("")).get()
 
-        func(unannotated_user) # E: Argument 1 to "func" has incompatible type "User"; expected "WithAnnotations[myapp.models.User, TypedDict('main.MyDict', {'foo': builtins.str})]"
+        func(unannotated_user) # E: Argument 1 to "func" has incompatible type "User"; expected "WithAnnotations[myapp__models__User, TypedDict('main.MyDict', {'foo': builtins.str})]"
         x: WithAnnotations[User]
         func(x)
         func(annotated_user)
-        func(other_annotated_user) # E: Argument 1 to "func" has incompatible type "WithAnnotations[myapp.models.User, TypedDict({'other': Any})]"; expected "WithAnnotations[myapp.models.User, TypedDict('main.MyDict', {'foo': builtins.str})]"
+        func(other_annotated_user) # E: Argument 1 to "func" has incompatible type "WithAnnotations[myapp__models__User, TypedDict({'other': Any})]"; expected "WithAnnotations[myapp__models__User, TypedDict('main.MyDict', {'foo': builtins.str})]"
 
         def func2(user: WithAnnotations[User, MyDict]) -> str:
-            print(user.asdf) # E: "WithAnnotations[myapp.models.User, TypedDict('main.MyDict', {'foo': builtins.str})]" has no attribute "asdf"
+            print(user.asdf) # E: "WithAnnotations[myapp__models__User, TypedDict('main.MyDict', {'foo': builtins.str})]" has no attribute "asdf"
             return user.foo
 
-        func2(unannotated_user) # E: Argument 1 to "func2" has incompatible type "User"; expected "WithAnnotations[myapp.models.User, TypedDict('main.MyDict', {'foo': builtins.str})]"
+        func2(unannotated_user) # E: Argument 1 to "func2" has incompatible type "User"; expected "WithAnnotations[myapp__models__User, TypedDict('main.MyDict', {'foo': builtins.str})]"
         func2(annotated_user)
-        func2(other_annotated_user) # E: Argument 1 to "func2" has incompatible type "WithAnnotations[myapp.models.User, TypedDict({'other': Any})]"; expected "WithAnnotations[myapp.models.User, TypedDict('main.MyDict', {'foo': builtins.str})]"
+        func2(other_annotated_user) # E: Argument 1 to "func2" has incompatible type "WithAnnotations[myapp__models__User, TypedDict({'other': Any})]"; expected "WithAnnotations[myapp__models__User, TypedDict('main.MyDict', {'foo': builtins.str})]"
     installed_apps:
         - myapp
     files:
@@ -100,7 +100,7 @@
         func(y)
 
         z: WithAnnotations[User, OtherDict]
-        func(z) # E: Argument 1 to "func" has incompatible type "WithAnnotations[myapp.models.User, TypedDict('main.OtherDict', {'other': builtins.str})]"; expected "WithAnnotations[myapp.models.User, TypedDict('main.NarrowDict', {'foo': builtins.str})]"
+        func(z) # E: Argument 1 to "func" has incompatible type "WithAnnotations[myapp__models__User, TypedDict('main.OtherDict', {'other': builtins.str})]"; expected "WithAnnotations[myapp__models__User, TypedDict('main.NarrowDict', {'foo': builtins.str})]"
 
     installed_apps:
         - myapp
@@ -119,12 +119,12 @@
         from django.db.models.expressions import F
 
         qs = User.objects.annotate(foo=F('id'))
-        reveal_type(qs) # N: Revealed type is "django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[myapp.models.User, TypedDict({'foo': Any})], django_stubs_ext.WithAnnotations[myapp.models.User, TypedDict({'foo': Any})]]"
+        reveal_type(qs) # N: Revealed type is "django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[myapp__models__User, TypedDict({'foo': Any})], django_stubs_ext.WithAnnotations[myapp__models__User, TypedDict({'foo': Any})]]"
 
         annotated = qs.get()
-        reveal_type(annotated) # N: Revealed type is "django_stubs_ext.WithAnnotations[myapp.models.User, TypedDict({'foo': Any})]*"
+        reveal_type(annotated) # N: Revealed type is "django_stubs_ext.WithAnnotations[myapp__models__User, TypedDict({'foo': Any})]*"
         reveal_type(annotated.foo) # N: Revealed type is "Any"
-        print(annotated.bar) # E: "WithAnnotations[myapp.models.User, TypedDict({'foo': Any})]" has no attribute "bar"
+        print(annotated.bar) # E: "WithAnnotations[myapp__models__User, TypedDict({'foo': Any})]" has no attribute "bar"
         reveal_type(annotated.username) # N: Revealed type is "builtins.str*"
 
     installed_apps:
@@ -144,7 +144,7 @@
         from django.db.models import Count
 
         qs = User.objects.annotate(Count('id'))
-        reveal_type(qs) # N: Revealed type is "django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[myapp.models.User], django_stubs_ext.WithAnnotations[myapp.models.User]]"
+        reveal_type(qs) # N: Revealed type is "django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[myapp__models__User], django_stubs_ext.WithAnnotations[myapp__models__User]]"
 
     installed_apps:
         - myapp
@@ -167,7 +167,7 @@
         def animals_only(param: Animal):
             pass
         # Make sure that even though attr access falls back to Any, the type is still checked
-        animals_only(annotated_user) # E: Argument 1 to "animals_only" has incompatible type "WithAnnotations[myapp.models.User]"; expected "Animal"
+        animals_only(annotated_user) # E: Argument 1 to "animals_only" has incompatible type "WithAnnotations[myapp__models__User]"; expected "Animal"
 
         def users_allowed(param: User):
             # But this function accepts only the original User type, so any attr access is not allowed within this function
@@ -196,7 +196,7 @@
         qs = User.objects.annotate(foo=F('id'))
         qs = qs.annotate(bar=F('id'))
         annotated = qs.get()
-        reveal_type(annotated) # N: Revealed type is "django_stubs_ext.WithAnnotations[myapp.models.User, TypedDict({'foo': Any, 'bar': Any})]*"
+        reveal_type(annotated) # N: Revealed type is "django_stubs_ext.WithAnnotations[myapp__models__User, TypedDict({'foo': Any, 'bar': Any})]*"
         reveal_type(annotated.foo) # N: Revealed type is "Any"
         reveal_type(annotated.bar) # N: Revealed type is "Any"
         reveal_type(annotated.username) # N: Revealed type is "builtins.str*"
@@ -227,11 +227,11 @@
           return qs.annotate(foo=F('id'))
 
         def add_wrong_annotation(qs: QuerySet[User]) -> QuerySet[WithAnnotations[User, FooDict]]:
-          return qs.annotate(bar=F('id')) # E: Incompatible return value type (got "_QuerySet[WithAnnotations[myapp.models.User, TypedDict({'bar': Any})], WithAnnotations[myapp.models.User, TypedDict({'bar': Any})]]", expected "_QuerySet[WithAnnotations[myapp.models.User, TypedDict('main.FooDict', {'foo': builtins.str})], WithAnnotations[myapp.models.User, TypedDict('main.FooDict', {'foo': builtins.str})]]")
+          return qs.annotate(bar=F('id')) # E: Incompatible return value type (got "_QuerySet[WithAnnotations[myapp__models__User, TypedDict({'bar': Any})], WithAnnotations[myapp__models__User, TypedDict({'bar': Any})]]", expected "_QuerySet[WithAnnotations[myapp__models__User, TypedDict('main.FooDict', {'foo': builtins.str})], WithAnnotations[myapp__models__User, TypedDict('main.FooDict', {'foo': builtins.str})]]")
 
         qs = add_annotation(qs)
         qs.get().foo
-        qs.get().bar # E: "WithAnnotations[myapp.models.User, TypedDict('main.FooDict', {'foo': builtins.str})]" has no attribute "bar"
+        qs.get().bar # E: "WithAnnotations[myapp__models__User, TypedDict('main.FooDict', {'foo': builtins.str})]" has no attribute "bar"
 
     installed_apps:
         - myapp