Repository: typeddjango/django-stubs

Issue:
Non-deterministic crash "KeyError: 'model_bases'" with auth models
# Bug report

I'm able to reproduce this crash mostly with a *clean* cache (`--no-incremental` or `rm -rf .mypy_cache`). If I have a warm cache from scanning some other files, this does not occur.

Just a small annoyance, I don't have time to debug right now, but figured I should report it.

```
.../site-packages/django-stubs/contrib/auth/models.pyi:21: error: INTERNAL ERROR -- Please try using mypy master on GitHub:
https://mypy.readthedocs.io/en/stable/common_issues.html#using-a-development-mypy-build
Please report a bug at https://github.com/python/mypy/issues
version: 1.0.1
Traceback (most recent call last):
  File "mypy/semanal.py", line 6210, in accept
  File "mypy/nodes.py", line 1133, in accept
  File "mypy/semanal.py", line 1545, in visit_class_def
  File "mypy/semanal.py", line 1625, in analyze_class
  File "mypy/semanal.py", line 1652, in analyze_class_body_common
  File "mypy/semanal.py", line 1735, in apply_class_plugin_hooks
  File ".../site-packages/mypy_django_plugin/main.py", line 44, in transform_model_class
    helpers.get_django_metadata(sym.node)["model_bases"][ctx.cls.fullname] = 1
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^
KeyError: 'model_bases'
.../site-packages/django-stubs/contrib/auth/models.pyi:21: : note: use --pdb to drop into pdb
```

## System information

- OS: macOS
- `python` version: 3.11.2
- `django` version: 4.1.7
- `mypy` version: 1.0.1
- `django-stubs` version: 1.15.0
- `django-stubs-ext` version: 0.7.0



Discussion:
I've just started experiencing this
@jonodrew Any idea what changed about your project that may have provoked this?

I haven't experienced this crash in my usage any more.

Suddenly now this happened to me as well. 🤷

The code is:

```python
def transform_model_class(ctx: ClassDefContext, django_context: DjangoContext) -> None:
    ...
    helpers.get_django_metadata(sym.node)["model_bases"][ctx.cls.fullname] = 1
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^

def get_django_metadata(model_info: TypeInfo) -> Dict[str, Any]:
    return model_info.metadata.setdefault("django", {})
```

So it should be entirely expected that `get_django_metadata()` can return empty dict.

Patching it like this at least avoids the exception though probably not correct:
```python
    return model_info.metadata.setdefault("django", {"model_bases": {}})
```

There's also another similar assumption made in another function `transform_form_class()`

The `"model_bases"` mapping is set by `_get_current_model_bases`.

https://github.com/typeddjango/django-stubs/blob/91dc6d31b8cb8ce49a985727331194ed26bb11c1/mypy_django_plugin/main.py#L94-L101

While the base class hooks e.g. `transform_model_class` is called depending on an `or` conditioned if-statement:

https://github.com/typeddjango/django-stubs/blob/91dc6d31b8cb8ce49a985727331194ed26bb11c1/mypy_django_plugin/main.py#L255-L263

I would suspect that `fullname in self.django_context.all_registered_model_class_fullnames` returns true at a point where `self._get_current_model_bases()` hasn't yet been called and the `or`-statement won't trigger it either.

As such `"model_bases"` hasn't yet been defaultly inserted. But this is just some quick and early guess work, and I have no way of reproducing. I just discovered the code paths.

We have multiple of these state mappings:

https://github.com/typeddjango/django-stubs/blob/91dc6d31b8cb8ce49a985727331194ed26bb11c1/mypy_django_plugin/main.py#L76-L115

Best way, IMO, would be to organise those data structures behind a proper interface. While the quick and dirty way could be to trigger the population inside `django_metadata` once at some early boot(?) point of the plugin

Resolving patch:
diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py
--- a/mypy_django_plugin/lib/helpers.py
+++ b/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
diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py
--- a/mypy_django_plugin/main.py
+++ b/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 {}
 


Test patch: