Wagtail Draftail code block

wagtail_hooks.py

import wagtail.admin.rich_text.editors.draftail.features as draftail_features
from wagtail.admin.rich_text.converters.html_to_contentstate import (
    BlockElementHandler, InlineStyleElementHandler
)
from wagtail.core import hooks


@hooks.register('register_rich_text_features')
def register_code_block_feature(features):
    """
    Registering the `code-block` feature, which uses the `code-block` Draft.js block type,
    and is stored as HTML with `<pre><code>` tags.
    """
    feature_name = 'code-block'
    type_ = 'code-block'

    control = {
        'type': type_,
        'label': '{}',
        'description': 'Code',
    }

    features.register_editor_plugin(
        'draftail', feature_name, draftail_features.BlockFeature(control)
    )

    features.register_converter_rule('contentstate', feature_name, {
        'from_database_format': {
            'pre': BlockElementHandler(type_),
            'code': InlineStyleElementHandler('CODE'),
        },
        'to_database_format': {
            'block_map': {'code-block': {'element': 'code', 'wrapper': 'pre'}}
        },
    })

    features.default_features.append(feature_name)