From 7e346e3278b21eb9d7601c55f8e86696f665b4b6 Mon Sep 17 00:00:00 2001 From: Bryce Seager van Dyk Date: Mon, 19 Feb 2024 11:27:59 +1300 Subject: [PATCH] Avoid returning None from get_field_related_model_cls https://github.com/typeddjango/django-stubs/pull/1495 updated `get_field_related_model_cls` to raise `UnregisteredModelError` rather than returning `None` for failure paths. However, None can still be returned if the initial retrieval of `related_model_cls` returns None. This patch adds a check to see if the initial retrieval has got a `None` and then raises the appropriate error rather than letting `None` be returned. --- mypy_django_plugin/django/context.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py index 347ec92ef5..43e5763219 100644 --- a/mypy_django_plugin/django/context.py +++ b/mypy_django_plugin/django/context.py @@ -1,7 +1,7 @@ import os import sys from collections import defaultdict -from contextlib import contextmanager +from contextlib import contextmanager, suppress from functools import cached_property from typing import ( TYPE_CHECKING, @@ -374,16 +374,14 @@ def get_field_related_model_cls(self, field: Union["RelatedField[Any, Any]", For # same model related_model_cls = field.model elif "." not in related_model_cls: - # same file model related_model_fullname = f"{field.model.__module__}.{related_model_cls}" related_model_cls = self.get_model_class_by_fullname(related_model_fullname) - if related_model_cls is None: - raise UnregisteredModelError else: - try: + with suppress(LookupError): related_model_cls = self.apps_registry.get_model(related_model_cls) - except LookupError as e: - raise UnregisteredModelError from e + + if related_model_cls is None: + raise UnregisteredModelError return related_model_cls