Fix plugin exception "KeyError: 'model_bases'" and related errors (#1563)

Author: intgr

mypy_django_plugin/lib/helpers.py

@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Set, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Literal, Optional, Set, Union, cast
 
 from django.db.models.fields import Field
 from django.db.models.fields.related import RelatedField
@@ -35,6 +35,7 @@
 from mypy.semanal import SemanticAnalyzer
 from mypy.types import AnyType, Instance, NoneTyp, TupleType, TypedDictType, TypeOfAny, UnionType
 from mypy.types import Type as MypyType
+from typing_extensions import TypedDict
 
 from mypy_django_plugin.lib import fullnames
 from mypy_django_plugin.lib.fullnames import WITH_ANNOTATIONS_FULLNAME
@@ -43,8 +44,23 @@
     from mypy_django_plugin.django.context import DjangoContext
 
 
-def get_django_metadata(model_info: TypeInfo) -> Dict[str, Any]:
-    return model_info.metadata.setdefault("django", {})
+class DjangoTypeMetadata(TypedDict, total=False):
+    from_queryset_manager: str
+    reverse_managers: Dict[str, str]
+    baseform_bases: Dict[str, int]
+    manager_bases: Dict[str, int]
+    model_bases: Dict[str, int]
+    queryset_bases: Dict[str, int]
+
+
+def get_django_metadata(model_info: TypeInfo) -> DjangoTypeMetadata:
+    return cast(DjangoTypeMetadata, model_info.metadata.setdefault("django", {}))
+
+
+def get_django_metadata_bases(
+    model_info: TypeInfo, key: Literal["baseform_bases", "manager_bases", "model_bases", "queryset_bases"]
+) -> Dict[str, int]:
+    return get_django_metadata(model_info).setdefault(key, cast(Dict[str, int], {}))
 
 
 class IncompleteDefnException(Exception):
@@ -376,4 +392,5 @@ def add_new_sym_for_info(info: TypeInfo, *, name: str, sym_type: MypyType, no_se
 def add_new_manager_base(api: SemanticAnalyzerPluginInterface, fullname: str) -> None:
     sym = api.lookup_fully_qualified_or_none(fullnames.MANAGER_CLASS_FULLNAME)
     if sym is not None and isinstance(sym.node, TypeInfo):
-        get_django_metadata(sym.node)["manager_bases"][fullname] = 1
+        bases = get_django_metadata_bases(sym.node, "manager_bases")
+        bases[fullname] = 1

mypy_django_plugin/main.py

@@ -43,7 +43,8 @@ def transform_model_class(ctx: ClassDefContext, django_context: DjangoContext) -
     sym = ctx.api.lookup_fully_qualified_or_none(fullnames.MODEL_CLASS_FULLNAME)
 
     if sym is not None and isinstance(sym.node, TypeInfo):
-        helpers.get_django_metadata(sym.node)["model_bases"][ctx.cls.fullname] = 1
+        bases = helpers.get_django_metadata_bases(sym.node, "model_bases")
+        bases[ctx.cls.fullname] = 1
     else:
         if not ctx.api.final_iteration:
             ctx.api.defer()
@@ -55,7 +56,8 @@ def transform_model_class(ctx: ClassDefContext, django_context: DjangoContext) -
 def transform_form_class(ctx: ClassDefContext) -> None:
     sym = ctx.api.lookup_fully_qualified_or_none(fullnames.BASEFORM_CLASS_FULLNAME)
     if sym is not None and isinstance(sym.node, TypeInfo):
-        helpers.get_django_metadata(sym.node)["baseform_bases"][ctx.cls.fullname] = 1
+        bases = helpers.get_django_metadata_bases(sym.node, "baseform_bases")
+        bases[ctx.cls.fullname] = 1
 
     forms.make_meta_nested_class_inherit_from_any(ctx)
 
@@ -77,41 +79,38 @@ def __init__(self, options: Options) -> None:
     def _get_current_queryset_bases(self) -> Dict[str, int]:
         model_sym = self.lookup_fully_qualified(fullnames.QUERYSET_CLASS_FULLNAME)
         if model_sym is not None and isinstance(model_sym.node, TypeInfo):
-            return helpers.get_django_metadata(model_sym.node).setdefault(  # type: ignore[no-any-return]
-                "queryset_bases", {fullnames.QUERYSET_CLASS_FULLNAME: 1}
-            )
+            bases = helpers.get_django_metadata_bases(model_sym.node, "queryset_bases")
+            bases[fullnames.QUERYSET_CLASS_FULLNAME] = 1
+            return bases
         else:
             return {}
 
     def _get_current_manager_bases(self) -> Dict[str, int]:
         model_sym = self.lookup_fully_qualified(fullnames.MANAGER_CLASS_FULLNAME)
         if model_sym is not None and isinstance(model_sym.node, TypeInfo):
-            return helpers.get_django_metadata(model_sym.node).setdefault(  # type: ignore[no-any-return]
-                "manager_bases", {fullnames.MANAGER_CLASS_FULLNAME: 1}
-            )
+            bases = helpers.get_django_metadata_bases(model_sym.node, "manager_bases")
+            bases[fullnames.MANAGER_CLASS_FULLNAME] = 1
+            return bases
         else:
             return {}
 
     def _get_current_model_bases(self) -> Dict[str, int]:
         model_sym = self.lookup_fully_qualified(fullnames.MODEL_CLASS_FULLNAME)
         if model_sym is not None and isinstance(model_sym.node, TypeInfo):
-            return helpers.get_django_metadata(model_sym.node).setdefault(  # type: ignore[no-any-return]
-                "model_bases", {fullnames.MODEL_CLASS_FULLNAME: 1}
-            )
+            bases = helpers.get_django_metadata_bases(model_sym.node, "model_bases")
+            bases[fullnames.MODEL_CLASS_FULLNAME] = 1
+            return bases
         else:
             return {}
 
     def _get_current_form_bases(self) -> Dict[str, int]:
         model_sym = self.lookup_fully_qualified(fullnames.BASEFORM_CLASS_FULLNAME)
         if model_sym is not None and isinstance(model_sym.node, TypeInfo):
-            return helpers.get_django_metadata(model_sym.node).setdefault(  # type: ignore[no-any-return]
-                "baseform_bases",
-                {
-                    fullnames.BASEFORM_CLASS_FULLNAME: 1,
-                    fullnames.FORM_CLASS_FULLNAME: 1,
-                    fullnames.MODELFORM_CLASS_FULLNAME: 1,
-                },
-            )
+            bases = helpers.get_django_metadata_bases(model_sym.node, "baseform_bases")
+            bases[fullnames.BASEFORM_CLASS_FULLNAME] = 1
+            bases[fullnames.FORM_CLASS_FULLNAME] = 1
+            bases[fullnames.MODELFORM_CLASS_FULLNAME] = 1
+            return bases
         else:
             return {}