Custom image entity in Draftail

CaptionedImageModalWorkflowSource.js

const basePlugins = window.draftail.registerPlugin({ type: 'potato '});
const imagePlugin = basePlugins.IMAGE;
const ImageModalWorkflowSource = imagePlugin.source; // ImageModalWorkflowSource
const imagePluginBlock = imagePlugin.block; // ImageBlock

class CaptionedImageModalWorkflowSource extends ImageModalWorkflowSource {
    getChooserConfig(entity) {
      let url;
      let urlParams;
  
      if (entity) {
        const data = entity.getData();
        url = `${global.chooserUrls.imageChooser}${data.id}/select_format/`;
        urlParams = {
          format: data.format,
          alt_text: data.alt,
            //   TODO: add caption
            caption: data.caption,
        };
      } else {
        url = `${global.chooserUrls.imageChooser}?select_format=true`;
        urlParams = {};
      }
      return {
        url,
        urlParams,
        onload: global.IMAGE_CHOOSER_MODAL_ONLOAD_HANDLERS,
        responses: {
          chosen: this.onChosen,
        },
      };
    }
  
    filterEntityData(data) {
      return {
        id: data.id,
        src: data.preview.url,
        alt: data.alt,
        format: data.format,
        caption: data.caption,
      };
    }
  }

window.draftail.registerPlugin({
    type: 'IMAGE',
    source: CaptionedImageModalWorkflowSource,
    block: ImageBlock,
});

wagtail_hooks.py

class ImageElementHandler(AtomicBlockEntityElementHandler):
    """
    Rule for building an image entity when converting from database representation
    to contentstate
    """

    def create_entity(self, name, attrs, state, contentstate):
        Image = get_image_model()
        try:
            image = Image.objects.get(id=attrs["id"])
            image_format = get_image_format(attrs["format"])
            rendition = get_rendition_or_not_found(image, image_format.filter_spec)
            src = rendition.url
        except Image.DoesNotExist:
            src = ""

        return Entity(
            "IMAGE",
            "IMMUTABLE",
            {
                "id": attrs["id"],
                "src": src,
                "alt": attrs.get("alt"),
                "format": attrs["format"],
                "caption": attrs.get("caption"),
            },
        )
        
def image_entity(props):
    """
    Helper to construct elements of the form
    <embed alt="Right-aligned image" embedtype="image" format="right" id="1"/>
    when converting from contentstate data
    """
    return DOM.create_element(
        "embed",
        {
            "embedtype": "image",
            "format": props.get("format"),
            "id": props.get("id"),
            "alt": props.get("alt"),
            "caption": props.get("caption"),
        },
    )


# <embed embedtype="image" id="10" alt="A pied wagtail" format="left" caption="Potato" />
ContentstateImageConversionRule = {
    "from_database_format": {
        'embed[embedtype="image"]': ImageElementHandler(),
    },
    "to_database_format": {"entity_decorators": {"IMAGE": image_entity}},
}

@hooks.register("register_rich_text_features")
def register_image_feature(features):
    # define a handler for converting <embed embedtype="image"> tags into frontend HTML
    features.register_embed_type(ImageEmbedHandler)

    # define a draftail plugin to use when the 'image' feature is active
    features.register_editor_plugin(
        "draftail",
        "image",
        draftail_features.EntityFeature(
            {
                "type": "IMAGE",
                "icon": "image",
                "description": gettext("Image"),
                # We do not want users to be able to copy-paste hotlinked images into rich text.
                # Keep only the attributes Wagtail needs.
                "attributes": ["id", "src", "alt", "format", "caption"],
                # Keep only images which are from Wagtail.
                "allowlist": {
                    "id": True,
                },
            },
            js=[
                "wagtailimages/js/image-chooser-modal.js",
            ],
        ),
    )

    # define how to convert between contentstate's representation of images and
    # the database representation
    features.register_converter_rule(
        "contentstate", "image", ContentstateImageConversionRule
    )