Split the _indexed_models dict into two sets, one for tracked models, the other for indexed models

Also tweak register_model to recursively follow child relations.

wagtail/models/reference_index.py

@@ -155,11 +155,15 @@ class ReferenceIndex(models.Model):
 
     wagtail_reference_index_ignore = True
 
-    # Store the models that are being tracked and indexed.
-    # If a model is present as a key, the model is tracked by receiving the corresponding
-    # model save and delete signals. If the value is True, then it is indexed.
-    # (False indicates the model has a ParentalKey relation with an indexed model.)
-    _indexed_models = {}
+    # The set of models that should have signals attached to watch for outbound references.
+    # This includes those registered with `register_model`, as well as their child models
+    # linked by a ParentalKey.
+    tracked_models = set()
+
+    # Ths set of models that can appear as the 'from' object in the reference index.
+    # This only includes those registered with `register_model`, and NOT child models linked
+    # by ParentalKey (object references on those are recorded under the parent).
+    indexed_models = set()
 
     class Meta:
         unique_together = [
@@ -245,29 +249,28 @@ class ReferenceIndex(models.Model):
     def register_model(cls, model):
         """
         Registers the model for indexing.
-
-        If there are child relationships (via a ParentalKey), the
-        saves and deletes on those models also need to be tracked
         """
         if cls.model_is_indexable(model):
-            cls._indexed_models[model] = True
+            cls.indexed_models.add(model)
+            cls._register_as_tracked_model(model)
 
-            for child_relation in get_all_child_relations(model):
-                if cls.model_is_indexable(
-                    child_relation.related_model,
-                    allow_child_models=True,
-                ):
-                    cls._indexed_models[
-                        child_relation.related_model
-                    ] = cls._indexed_models.get(child_relation.related_model, False)
+    @classmethod
+    def _register_as_tracked_model(cls, model):
+        """
+        Add the model and all of its ParentalKey-linked children to the set of
+        models to be tracked by signal handlers.
+        """
+        cls.tracked_models.add(model)
+        for child_relation in get_all_child_relations(model):
+            if cls.model_is_indexable(
+                child_relation.related_model,
+                allow_child_models=True,
+            ):
+                cls._register_as_tracked_model(child_relation.related_model)
 
     @classmethod
     def is_indexed(cls, model):
-        return cls._indexed_models.get(model, False)
-
-    @classmethod
-    def get_tracked_models(cls):
-        return cls._indexed_models.keys()
+        return model in cls.indexed_models
 
     @classmethod
     def _extract_references_from_object(cls, object):

wagtail/signal_handlers.py

@@ -100,13 +100,13 @@ def remove_reference_index_on_delete(instance, **kwargs):
 
 
 def connect_reference_index_signal_handlers(**kwargs):
-    for model in ReferenceIndex.get_tracked_models():
+    for model in ReferenceIndex.tracked_models:
         post_save.connect(update_reference_index_on_save, sender=model)
         post_delete.connect(remove_reference_index_on_delete, sender=model)
 
 
 def disconnect_reference_index_signal_handlers(**kwargs):
-    for model in ReferenceIndex.get_tracked_models():
+    for model in ReferenceIndex.tracked_models:
         post_save.disconnect(update_reference_index_on_save, sender=model)
         post_delete.disconnect(remove_reference_index_on_delete, sender=model)