Snippet embed feature starting

wagtail_draftail_snippet/richtext.py

@@ -2,12 +2,15 @@ from django.apps import apps
 from django.template.loader import render_to_string
 
 from draftjs_exporter.dom import DOM
-from wagtail.admin.rich_text.converters.html_to_contentstate import LinkElementHandler
+from wagtail.admin.rich_text.converters.contentstate_models import Entity
-from wagtail.core.rich_text import LinkHandler
+from wagtail.admin.rich_text.converters.html_to_contentstate import LinkElementHandler, AtomicBlockEntityElementHandler
+from wagtail.core.rich_text import EmbedHandler, LinkHandler
 
-from .utils import get_snippet_frontend_template
+from .utils import get_snippet_frontend_template, get_snippet_embed_frontend_template
 
 
+# Snippet Link
+
 # Front-end conversion
 class SnippetLinkHandler(LinkHandler):
     identifier = "snippet"
@@ -82,3 +85,77 @@ ContentstateSnippetLinkConversionRule = {
     },
     "to_database_format": {"entity_decorators": {"SNIPPET": snippet_link_entity}},
 }
+
+
+# Snippet Embed
+
+# Front-end conversion
+class SnippetEmbedHandler(EmbedHandler):
+    identifier = "snippet"
+
+    @classmethod
+    def get_instance(cls, attrs):
+        model = apps.get_model(attrs["data-app-name"], attrs["data-model-name"])
+        return model.objects.get(id=attrs["id"])
+
+    @classmethod
+    def get_template(cls, attrs):
+        return get_snippet_embed_frontend_template(
+            attrs["data-app-name"], attrs["data-model-name"]
+        )
+
+    @classmethod
+    def expand_db_attributes(cls, attrs):
+        try:
+            snippet_obj = cls.get_instance(attrs)
+            template = cls.get_template(attrs)
+            return render_to_string(template, {"object": snippet_obj})
+        except Exception:
+            return ""
+
+
+# draft.js / contentstate conversion
+def snippet_embed_entity(props):
+    """
+    Helper to construct elements of the form
+    <embed embedtype="snippet" id="1"/> when converting from contentstate data
+    """
+
+    elem = DOM.create_element(
+        "embed",
+        {
+            "embedtype": "snippet",
+            "id": props.get("id"),
+            "data-string": props.get("string"),
+            "data-edit-link": props.get("edit_link"),
+            "data-app-name": props.get("app_name"),
+            "data-model-name": props.get("model_name"),
+        }
+    )
+
+    return elem
+
+
+class SnippetEmbedElementHandler(AtomicBlockEntityElementHandler):
+    """
+    Rule for building a snippet entity when converting from database representation
+    to contentstate
+    """
+
+    def create_entity(self, name, attrs, state, contentstate):
+        return Entity('SNIPPET-EMBED', 'IMMUTABLE', {
+            "embedtype": "snippet",
+            "id": attrs.get("id"),
+            "data-string": attrs.get("string"),
+            "data-edit-link": attrs.get("edit_link"),
+            "data-app-name": attrs.get("app_name"),
+            "data-model-name": attrs.get("model_name"),
+        })
+
+
+ContentstateSnippetEmbedConversionRule = {
+    "from_database_format": {
+        'embed[embedtype="snippet"]': SnippetEmbedElementHandler()
+    },
+    "to_database_format": {"entity_decorators": {"SNIPPET-EMBED": snippet_embed_entity}},
+}

wagtail_draftail_snippet/static/wagtail_draftail_snippet/js/wagtail_draftail_snippet.js

@@ -11,6 +11,10 @@
 
   const $ = global.jQuery;
 
+  const MUTABILITY = {};
+  MUTABILITY['SNIPPET'] = 'MUTABLE';
+  MUTABILITY['SNIPPET-EMBED'] = 'IMMUTABLE';
+
   const getSnippetModelChooserConfig = () => {
     let url;
     let urlParams;
@@ -113,7 +117,7 @@
       const selection = editorState.getSelection();
 
       const entityData = filterSnippetEntityData(entityType, data);
-      const mutability = 'MUTABLE';
+      const mutability = MUTABILITY[entityType.type];
       const contentWithEntity = content.createEntity(entityType.type, mutability, entityData);
       const entityKey = contentWithEntity.getLastCreatedEntityKey();
 
@@ -177,9 +181,32 @@
     });
   };
 
+  const SnippetEmbed = props => {
+    const { entity, onRemoveEntity } = props.blockProps;
+    const data = entity.getData();
+
+    let icon = React.createElement(window.wagtail.components.Icon, {name: 'media'});
+    let label = data.string || '';
+
+    return React.createElement(TooltipEntity, {
+      entityKey: props.entityKey,
+      children: props.children,
+      onEdit: props.onEdit,
+      onRemove: props.onRemove,
+      icon: icon,
+      label: label
+    });
+  };
+
   window.draftail.registerPlugin({
     type: 'SNIPPET',
     source: SnippetModalWorkflowSource,
     decorator: SnippetLink,
   });
+
+  window.draftail.registerPlugin({
+    type: 'SNIPPET-EMBED',
+    source: SnippetModalWorkflowSource,
+    block: SnippetEmbed,
+  });
 })();

wagtail_draftail_snippet/utils.py

@@ -3,3 +3,7 @@ from wagtail.core.utils import camelcase_to_underscore
 
 def get_snippet_frontend_template(app_name, model_name):
     return "%s/%s_snippet.html" % (app_name, camelcase_to_underscore(model_name))
+
+
+def get_snippet_embed_frontend_template(app_name, model_name):
+    return "%s/%s_snippet_embed.html" % (app_name, camelcase_to_underscore(model_name))

wagtail_draftail_snippet/views.py

@@ -3,7 +3,7 @@ from django.template.loader import TemplateDoesNotExist, get_template
 from wagtail.admin.modal_workflow import render_modal_workflow
 from wagtail.snippets.models import get_snippet_models
 
-from .utils import get_snippet_frontend_template
+from .utils import get_snippet_frontend_template, get_snippet_embed_frontend_template
 
 
 def choose_snippet_model(request):
@@ -11,15 +11,29 @@ def choose_snippet_model(request):
 
     # Only display those snippet models which have snippet frontend template
     for snippet_model in get_snippet_models():
+        snippet_included = False
         snippet_frontend_template = get_snippet_frontend_template(
             snippet_model._meta.app_label, snippet_model._meta.model_name
         )
+        snippet_embed_frontend_template = get_snippet_embed_frontend_template(
+            snippet_model._meta.app_label, snippet_model._meta.model_name
+        )
+
         try:
             get_template(snippet_frontend_template)
             snippet_model_opts.append(snippet_model._meta)
+            snippet_included = True
         except TemplateDoesNotExist:
             pass
 
+        if not snippet_included:
+            try:
+                get_template(snippet_embed_frontend_template)
+                snippet_model_opts.append(snippet_model._meta)
+                snippet_included = True
+            except TemplateDoesNotExist:
+                pass
+
     return render_modal_workflow(
         request,
         "wagtail_draftail_snippet/choose_snippet_model.html",

wagtail_draftail_snippet/wagtail_hooks.py

@@ -7,7 +7,10 @@ import wagtail.admin.rich_text.editors.draftail.features as draftail_features
 from wagtail.core import hooks
 
 from . import urls
-from .richtext import ContentstateSnippetLinkConversionRule, SnippetLinkHandler
+from .richtext import (
+    ContentstateSnippetLinkConversionRule, ContentstateSnippetEmbedConversionRule,
+    SnippetEmbedHandler, SnippetLinkHandler
+)
 
 
 @hooks.register("register_rich_text_features")
@@ -35,6 +38,31 @@ def register_snippet_feature(features):
     )
 
 
+@hooks.register("register_rich_text_features")
+def register_snippet_embed_feature(features):
+    feature_name = "snippet-embed"
+    type_ = "SNIPPET-EMBED"
+
+    features.register_embed_type(SnippetEmbedHandler)
+
+    features.register_editor_plugin(
+        "draftail",
+        feature_name,
+        draftail_features.EntityFeature(
+            {"type": type_, "icon": "media", "description": ugettext("Snippet Embed")},
+            js=[
+                "wagtailsnippets/js/snippet-chooser-modal.js",
+                "wagtail_draftail_snippet/js/snippet-model-chooser-modal.js",
+                "wagtail_draftail_snippet/js/wagtail_draftail_snippet.js",
+            ],
+        ),
+    )
+
+    features.register_converter_rule(
+        "contentstate", feature_name, ContentstateSnippetEmbedConversionRule
+    )
+
+
 @hooks.register("insert_editor_js")
 def editor_js():
     return format_html(