From 072c8f21e13f92cdeff29584bd8e4abce60b7ad7 Mon Sep 17 00:00:00 2001 From: Matt Westcott Date: Wed, 19 Apr 2023 19:36:14 +0100 Subject: [PATCH] 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 | 47 ++++++++++++++++--------------- wagtail/signal_handlers.py | 4 +-- 2 files changed, 27 insertions(+), 24 deletions(-) diff --git a/wagtail/models/reference_index.py b/wagtail/models/reference_index.py index fab73cc06f..80739ef414 100644 --- a/wagtail/models/reference_index.py +++ b/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): diff --git a/wagtail/signal_handlers.py b/wagtail/signal_handlers.py index 2f2f287906..7db2867dcb 100644 --- a/wagtail/signal_handlers.py +++ b/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)