from django.core.exceptions import ImproperlyConfigured
from django.db import models


class ObjectTypeRegistry:
    """
    Implements a lookup table for mapping objects to values according to the object type.
    The most specific type according to the object's inheritance chain is selected.
    """

    def __init__(self):
        # values in this dict will be returned if the field type exactly matches an item here
        self.values_by_exact_class = {}

        # values in this dict will be returned if any class in the field's inheritance chain
        # matches, preferring more specific subclasses
        self.values_by_class = {}

    def register(self, cls, value=None, exact_class=False):
        if exact_class:
            self.values_by_exact_class[cls] = value
        else:
            self.values_by_class[cls] = value

    def get(self, obj):
        value = None
        if obj.__class__ in self.values_by_exact_class:
            value = self.values_by_exact_class[obj.__class__]
        else:
            for klass in obj.__class__.mro():
                if klass in self.values_by_class:
                    value = self.values_by_class[klass]
                    break

        if callable(value) and not isinstance(value, type):
            value = value(obj)

        return value


class ModelFieldRegistry(ObjectTypeRegistry):
    """
    Handles the recurring pattern where we need to register different values for different
    model field types, and retrieve the one that most closely matches a given model field,
    according to its type (taking inheritance into account), and in the case of foreign keys,
    the type of the related model (again, taking inheritance into account).

    For example, this is used by wagtail.admin.forms.models when constructing model forms:
    we use such a registry to retrieve the appropriate dict of arguments to pass to the
    form field constructor. A lookup for a models.TextField will return a dict specifying a
    text area widget, and a lookup for a foreign key to Image will return a dict specifying
    an image chooser widget.
    """

    def __init__(self):
        super().__init__()
        self.values_by_class[models.ForeignKey] = self.foreign_key_lookup

        # values in this dict will be returned if the field is a foreign key to a related
        # model in here, matching most specific subclass first
        self.values_by_fk_related_model = {}

    def register(self, field_class, to=None, value=None, exact_class=False):
        if to:
            if field_class == models.ForeignKey:
                self.values_by_fk_related_model[to] = value
            else:
                raise ImproperlyConfigured(
                    "The 'to' argument on ModelFieldRegistry.register is only valid for ForeignKey fields"
                )
        else:
            super().register(field_class, value=value, exact_class=exact_class)

    def foreign_key_lookup(self, field):
        value = None
        target_model = field.remote_field.model

        for model in target_model.mro():
            if model in self.values_by_fk_related_model:
                value = self.values_by_fk_related_model[model]
                break

        if callable(value) and not isinstance(value, type):
            value = value(field)

        return value

@@ -851,18 +851,26 @@
  get
+_by_type
 (self, 
-obj
+cls
 ):%0A 
@@ -880,75 +880,11 @@
     
-value = None%0A        if obj.__class__ in self.values_by_exact_class
+try
 :%0A  
@@ -885,39 +885,38 @@
 ry:%0A            
-value =
+return
  self.values_by_
@@ -931,131 +931,96 @@
 ass%5B
-obj.__class__%5D%0A        else:%0A            for klass in obj.__class__.mro():%0A                if klass in self.values_by_class
+cls%5D%0A        except KeyError:%0A            for ancestor in cls.mro():%0A                try
 :%0A  
@@ -1029,39 +1029,38 @@
                 
-value =
+return
  self.values_by_
@@ -1069,13 +1069,16 @@
 ass%5B
-klass
+ancestor
 %5D%0A  
@@ -1087,33 +1087,138 @@
             
-    break
+except KeyError:%0A                    pass%0A%0A    def get(self, obj):%0A        value = self.get_by_type(obj.__class__)
 %0A%0A        if

Extract a get_by_type method