Register embed/link handlers through rich text features instead of hooks

expand_db_html simply uses all registered handlers at this point, since it's a global function that can't be customised through features

wagtail/admin/tests/test_rich_text.py

@@ -427,3 +427,35 @@ class TestWidgetWhitelisting(TestCase, WagtailTestUtils):
             'body': '<h1>h1</h1> <h2>h2</h2> <script>script</script> <p><b>bold</b> <i>italic</i></p> <blockquote>blockquote</blockquote>'
         }, {}, 'body')
         self.assertEqual(result, '<h1>h1</h1> h2 script <p><b>bold</b> italic</p> blockquote')
+
+    def test_link_conversion_with_default_whitelist(self):
+        widget = HalloRichTextArea()
+
+        result = widget.value_from_datadict({
+            'body': '<p>a <a href="/foo" data-linktype="page" data-id="123">page</a>, <a href="/foo" data-linktype="squirrel" data-id="234">a squirrel</a> and a <a href="/foo" data-linktype="document" data-id="345">document</a></p>'
+        }, {}, 'body')
+        self.assertHTMLEqual(result, '<p>a <a linktype="page" id="123">page</a>, a squirrel and a <a linktype="document" id="345">document</a></p>')
+
+    def test_link_conversion_with_custom_whitelist(self):
+        widget = HalloRichTextArea(features=['h1', 'bold', 'link', 'somethingijustmadeup'])
+
+        result = widget.value_from_datadict({
+            'body': '<p>a <a href="/foo" data-linktype="page" data-id="123">page</a>, <a href="/foo" data-linktype="squirrel" data-id="234">a squirrel</a> and a <a href="/foo" data-linktype="document" data-id="345">document</a></p>'
+        }, {}, 'body')
+        self.assertHTMLEqual(result, '<p>a <a linktype="page" id="123">page</a>, a squirrel and a document</p>')
+
+    def test_embed_conversion_with_default_whitelist(self):
+        widget = HalloRichTextArea()
+
+        result = widget.value_from_datadict({
+            'body': '<p>image <img src="foo" data-embedtype="image" data-id="123" data-format="left" data-alt="test alt" /> embed <span data-embedtype="media" data-url="https://www.youtube.com/watch?v=vwyuB8QKzBI">blah</span> badger <span data-embedtype="badger" data-colour="black-and-white">badger</span></p>'
+        }, {}, 'body')
+        self.assertHTMLEqual(result, '<p>image <embed embedtype="image" id="123" format="left" alt="test alt" /> embed <embed embedtype="media" url="https://www.youtube.com/watch?v=vwyuB8QKzBI" /> badger </p>')
+
+    def test_embed_conversion_with_custom_whitelist(self):
+        widget = HalloRichTextArea(features=['h1', 'bold', 'image', 'somethingijustmadeup'])
+
+        result = widget.value_from_datadict({
+            'body': '<p>image <img src="foo" data-embedtype="image" data-id="123" data-format="left" data-alt="test alt" /> embed <span data-embedtype="media" data-url="https://www.youtube.com/watch?v=vwyuB8QKzBI">blah</span></p>'
+        }, {}, 'body')
+        self.assertHTMLEqual(result, '<p>image <embed embedtype="image" id="123" format="left" alt="test alt" /> embed </p>')

wagtail/core/rich_text/__init__.py

@@ -3,7 +3,6 @@ from django.utils.safestring import mark_safe
 
 from wagtail.core import hooks
 from wagtail.core.rich_text.feature_registry import FeatureRegistry
-from wagtail.core.rich_text.pages import PageLinkHandler
 from wagtail.core.rich_text.rewriters import EmbedRewriter, LinkRewriter, MultiRuleRewriter
 from wagtail.core.whitelist import allow_without_attributes, Whitelister, DEFAULT_ELEMENT_RULES
 
@@ -51,21 +50,27 @@ class DbWhitelister(Whitelister):
 
     @cached_property
     def embed_handlers(self):
+        if self.features is None:
+            feature_list = features.get_default_features()
+        else:
+            feature_list = self.features
+
         embed_handlers = {}
-        for hook in hooks.get_hooks('register_rich_text_embed_handler'):
-            handler_name, handler = hook()
-            embed_handlers[handler_name] = handler
+        for feature in feature_list:
+            embed_handlers.update(features.get_embed_handler_rules(feature))
 
         return embed_handlers
 
     @cached_property
     def link_handlers(self):
-        link_handlers = {
-            'page': PageLinkHandler,
-        }
-        for hook in hooks.get_hooks('register_rich_text_link_handler'):
-            handler_name, handler = hook()
-            link_handlers[handler_name] = handler
+        if self.features is None:
+            feature_list = features.get_default_features()
+        else:
+            feature_list = self.features
+
+        link_handlers = {}
+        for feature in feature_list:
+            link_handlers.update(features.get_link_handler_rules(feature))
 
         return link_handlers
 
@@ -73,7 +78,13 @@ class DbWhitelister(Whitelister):
         if 'data-embedtype' in tag.attrs:
             embed_type = tag['data-embedtype']
             # fetch the appropriate embed handler for this embedtype
-            embed_handler = self.embed_handlers[embed_type]
+            try:
+                embed_handler = self.embed_handlers[embed_type]
+            except KeyError:
+                # discard embeds with unrecognised embedtypes
+                tag.decompose()
+                return
+
             embed_attrs = embed_handler.get_db_attributes(tag)
             embed_attrs['embedtype'] = embed_type
 
@@ -86,7 +97,13 @@ class DbWhitelister(Whitelister):
                 self.clean_node(doc, child)
 
             link_type = tag['data-linktype']
-            link_handler = self.link_handlers[link_type]
+            try:
+                link_handler = self.link_handlers[link_type]
+            except KeyError:
+                # discard links with unrecognised linktypes
+                tag.unwrap()
+                return
+
             link_attrs = link_handler.get_db_attributes(tag)
             link_attrs['linktype'] = link_type
             tag.attrs.clear()
@@ -99,8 +116,8 @@ class DbWhitelister(Whitelister):
 
 
 # Rewriter functions to be built up on first call to expand_db_html, using the utility classes
-# from wagtail.core.rich_text.rewriters along with the register_rich_text_embed_handler /
-# register_rich_text_link_handler hooks
+# from wagtail.core.rich_text.rewriters along with the embed handlers / link handlers registered
+# with the feature registry
 
 FRONTEND_REWRITER = None
 EDITOR_REWRITER = None
@@ -116,15 +133,16 @@ def expand_db_html(html, for_editor=False):
     if for_editor:
 
         if EDITOR_REWRITER is None:
-            embed_rules = {}
-            for hook in hooks.get_hooks('register_rich_text_embed_handler'):
-                handler_name, handler = hook()
-                embed_rules[handler_name] = handler.expand_db_attributes_for_editor
-
-            link_rules = {}
-            for hook in hooks.get_hooks('register_rich_text_link_handler'):
-                handler_name, handler = hook()
-                link_rules[handler_name] = handler.expand_db_attributes_for_editor
+            embed_handlers = features.get_all_embed_handler_rules()
+            embed_rules = {
+                handler_name: handler.expand_db_attributes_for_editor
+                for handler_name, handler in embed_handlers.items()
+            }
+            link_handlers = features.get_all_link_handler_rules()
+            link_rules = {
+                handler_name: handler.expand_db_attributes_for_editor
+                for handler_name, handler in link_handlers.items()
+            }
 
             EDITOR_REWRITER = MultiRuleRewriter([
                 LinkRewriter(link_rules), EmbedRewriter(embed_rules)
@@ -135,15 +153,16 @@ def expand_db_html(html, for_editor=False):
     else:
 
         if FRONTEND_REWRITER is None:
-            embed_rules = {}
-            for hook in hooks.get_hooks('register_rich_text_embed_handler'):
-                handler_name, handler = hook()
-                embed_rules[handler_name] = handler.expand_db_attributes
-
-            link_rules = {}
-            for hook in hooks.get_hooks('register_rich_text_link_handler'):
-                handler_name, handler = hook()
-                link_rules[handler_name] = handler.expand_db_attributes
+            embed_handlers = features.get_all_embed_handler_rules()
+            embed_rules = {
+                handler_name: handler.expand_db_attributes
+                for handler_name, handler in embed_handlers.items()
+            }
+            link_handlers = features.get_all_link_handler_rules()
+            link_rules = {
+                handler_name: handler.expand_db_attributes
+                for handler_name, handler in link_handlers.items()
+            }
 
             FRONTEND_REWRITER = MultiRuleRewriter([
                 LinkRewriter(link_rules), EmbedRewriter(embed_rules)

wagtail/core/rich_text/feature_registry.py

@@ -30,6 +30,14 @@ class FeatureRegistry:
         # the whitelister element_rules config when the feature is active
         self.whitelister_element_rules = {}
 
+        # a mapping of feature names to embed_handler rules that should be merged into the
+        # list of recognised embedtypes when the feature is active
+        self.embed_handler_rules = {}
+
+        # a mapping of feature names to link_handler rules that should be merged into the
+        # list of recognised linktypes when the feature is active
+        self.link_handler_rules = {}
+
     def get_default_features(self):
         if not self.has_scanned_for_features:
             self._scan_for_features()
@@ -62,3 +70,49 @@ class FeatureRegistry:
             self._scan_for_features()
 
         return self.whitelister_element_rules.get(feature_name, {})
+
+    def register_embed_handler_rules(self, feature_name, ruleset):
+        self.embed_handler_rules[feature_name] = ruleset
+
+    def get_embed_handler_rules(self, feature_name):
+        if not self.has_scanned_for_features:
+            self._scan_for_features()
+
+        return self.embed_handler_rules.get(feature_name, {})
+
+    def get_all_embed_handler_rules(self):
+        """
+        Return a dictionary of embedtypes to embed handlers, collated from all the
+        registered embed handler rules
+        """
+        if not self.has_scanned_for_features:
+            self._scan_for_features()
+
+        collated_ruleset = {}
+        for ruleset in self.embed_handler_rules.values():
+            collated_ruleset.update(ruleset)
+
+        return collated_ruleset
+
+    def register_link_handler_rules(self, feature_name, ruleset):
+        self.link_handler_rules[feature_name] = ruleset
+
+    def get_link_handler_rules(self, feature_name):
+        if not self.has_scanned_for_features:
+            self._scan_for_features()
+
+        return self.link_handler_rules.get(feature_name, {})
+
+    def get_all_link_handler_rules(self):
+        """
+        Return a dictionary of linktypes to link handlers, collated from all the
+        registered link handler rules
+        """
+        if not self.has_scanned_for_features:
+            self._scan_for_features()
+
+        collated_ruleset = {}
+        for ruleset in self.link_handler_rules.values():
+            collated_ruleset.update(ruleset)
+
+        return collated_ruleset

wagtail/core/wagtail_hooks.py

@@ -35,11 +35,6 @@ def check_view_restrictions(page, request, serve_args, serve_kwargs):
                 return require_wagtail_login(next=request.get_full_path())
 
 
-@hooks.register('register_rich_text_link_handler')
-def register_page_link_handler():
-    return ('page', PageLinkHandler)
-
-
 @hooks.register('register_rich_text_features')
 def register_core_features(features):
     features.default_features.append('hr')
@@ -47,6 +42,7 @@ def register_core_features(features):
 
     features.default_features.append('link')
     features.register_whitelister_element_rules('link', {'a': attribute_rule({'href': check_url})})
+    features.register_link_handler_rules('link', {'page': PageLinkHandler})
 
     features.default_features.append('bold')
     features.register_whitelister_element_rules(

wagtail/documents/wagtail_hooks.py

@@ -72,7 +72,7 @@ def editor_js():
 
 
 @hooks.register('register_rich_text_features')
-def register_embed_feature(features):
+def register_document_feature(features):
     features.register_editor_plugin(
         'hallo', 'document-link',
         HalloPlugin(
@@ -80,14 +80,10 @@ def register_embed_feature(features):
             js=['wagtaildocs/js/hallo-plugins/hallo-wagtaildoclink.js'],
         )
     )
+    features.register_link_handler_rules('document-link', {'document': DocumentLinkHandler})
     features.default_features.append('document-link')
 
 
-@hooks.register('register_rich_text_link_handler')
-def register_document_link_handler():
-    return ('document', DocumentLinkHandler)
-
-
 class DocumentsSummaryItem(SummaryItem):
     order = 300
     template = 'wagtaildocs/homepage/site_summary_documents.html'

wagtail/embeds/wagtail_hooks.py

@@ -36,9 +36,5 @@ def register_embed_feature(features):
             js=['wagtailembeds/js/hallo-plugins/hallo-wagtailembeds.js'],
         )
     )
+    features.register_embed_handler_rules('embed', {'media': MediaEmbedHandler})
     features.default_features.append('embed')
-
-
-@hooks.register('register_rich_text_embed_handler')
-def register_media_embed_handler():
-    return ('media', MediaEmbedHandler)

wagtail/images/wagtail_hooks.py

@@ -72,6 +72,7 @@ def register_image_feature(features):
             js=['wagtailimages/js/hallo-plugins/hallo-wagtailimage.js'],
         )
     )
+    features.register_embed_handler_rules('image', {'image': ImageEmbedHandler})
     features.default_features.append('image')
 
 
@@ -90,11 +91,6 @@ def register_image_operations():
     ]
 
 
-@hooks.register('register_rich_text_embed_handler')
-def register_image_embed_handler():
-    return ('image', ImageEmbedHandler)
-
-
 class ImagesSummaryItem(SummaryItem):
     order = 200
     template = 'wagtailimages/homepage/site_summary_images.html'