From f7a8e14b359b882a8d783996e028ae44c10de02d Mon Sep 17 00:00:00 2001
From: aaronhaslett <aaron.haslett@springload.co.n>
Date: Mon, 31 May 2021 10:37:07 +1200
Subject: [PATCH] Improvements to draftail option handling

---
 wagtailmodelchoosers/utils.py         |  5 +++++
 wagtailmodelchoosers/wagtail_hooks.py | 20 +++++++++++++-------
 wagtailmodelchoosers/widgets.py       | 10 ++++------
 3 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/wagtailmodelchoosers/utils.py b/wagtailmodelchoosers/utils.py
index 682faec..4fd383d 100644
--- a/wagtailmodelchoosers/utils.py
+++ b/wagtailmodelchoosers/utils.py
@@ -68,6 +68,11 @@ def get_chooser_options(chooser):
     return copy.deepcopy(chooser_options)
 
 
+def get_all_chooser_options():
+    choosers_options = getattr(settings, "MODEL_CHOOSERS_OPTIONS", {})
+    return {c: get_chooser_options(c) for c in choosers_options.keys()}
+
+
 def first_non_empty(data_or_instance, field_or_fields, default=None):
     """
     Return the first non-empty attribute of `instance` or `default`
diff --git a/wagtailmodelchoosers/wagtail_hooks.py b/wagtailmodelchoosers/wagtail_hooks.py
index 233ba40..f109ab8 100644
--- a/wagtailmodelchoosers/wagtail_hooks.py
+++ b/wagtailmodelchoosers/wagtail_hooks.py
@@ -6,6 +6,7 @@
 from wagtail.core import hooks
 
 from wagtailmodelchoosers.views import ModelView, RemoteResourceView
+from .utils import get_all_chooser_options
 
 
 @hooks.register("insert_editor_css")
@@ -48,12 +49,17 @@ def wagtailmodelchoosers_admin_urls():
 
 @hooks.register("register_rich_text_features")
 def register_rich_text_features(features):
-    entity_types = getattr(settings, "MODEL_CHOOSER_DRAFTAIL_ENTITY_TYPES", {})
-    for entity_type in entity_types:
-        type_ = entity_type.get("type", None)
-        if not type_:
+    choosers = get_all_chooser_options();
+    for name, chooser in choosers.items():
+        draftail_type = chooser.get('draftail_type', None)
+        if not draftail_type:
             continue
 
-        features.default_features.append(type_)
-        feature = draftail_features.EntityFeature(entity_type)
-        features.register_editor_plugin("draftail", type_, feature)
+        # If there's no content_type then it's a remote chooser, and
+        # content_type means the name of the chooser instead of the model.
+        if 'content_type' not in chooser:
+            chooser['content_type'] = name
+
+        chooser['type'] = draftail_type
+        feature = draftail_features.EntityFeature(chooser)
+        features.register_editor_plugin("draftail", name, feature)
diff --git a/wagtailmodelchoosers/widgets.py b/wagtailmodelchoosers/widgets.py
index d60f1cb..a5d3e81 100644
--- a/wagtailmodelchoosers/widgets.py
+++ b/wagtailmodelchoosers/widgets.py
@@ -10,7 +10,7 @@
 from wagtail.admin.staticfiles import versioned_static
 from wagtail.utils.widgets import WidgetWithScript
 
-from .utils import first_non_empty
+from .utils import first_non_empty, get_all_chooser_options
 
 
 class ModelChooserWidget(WidgetWithScript, widgets.Input):
@@ -222,9 +222,7 @@ def media(self):
 
     def get_context(self, *args, **kwargs):
         context = super().get_context(*args, **kwargs)
-        entity_types = getattr(settings, "MODEL_CHOOSER_DRAFTAIL_ENTITY_TYPES", {})
-        field = "type"
-        context["entity_types"] = json.dumps(
-            [e[field] for e in entity_types if field in e]
-        )
+        f = "draftail_type"
+        met = [c[f] for c in get_all_chooser_options().values() if f in c]
+        context["modelchooser_entity_types"] = json.dumps(met)
         return context