module to help migrate rich text content created by wagtaildraftail library to wagtail 2 presentation

wagtaildraftail_migration_tools.py

# -*- coding: utf-8 -*-
"""
module to help migrate rich text content created by wagtaildraftail library to wagtail 2 presentation
"""

from django.db import migrations
from wagtail.admin.rich_text.converters.contentstate import ContentstateConverter
from wagtail.core.blocks import StreamValue, StructValue
from wagtail.core.rich_text import features as feature_registry, RichText


def convert_stream_field_rich_text_blocks_migration(app, model, field_name):
    """
    Creates a django RunPython migration to migrate json content of DraftailTextBlock to html content of RichTextBlock
    inside StreamField recursive structure

    Usage:
    Create empty data migration and put line similar to this for each stream field that needs to be migrated
    ```
        operations = [
            convert_stream_field_rich_text_blocks_migration('core', 'BasePage', 'body'),
        ]
    ```
    PageModel = ContentType.objects.get_by_natural_key(app, model)
    :param app: model app name as in django content type
    :param model: model name as in django content type
    :param field_name: rich text field name to migrate values for
    :return: django.db.migrations.operations.special.RunPython
    """
    def migrate_rich_text_blocks(apps, schema_editor):
        PageModel = apps.get_model(app, model)
        convert_stream_field_rich_text_blocks(PageModel, field_name, convert_json_to_html)

    def reverse_migrate_rich_text_blocks(apps, schema_editor):
        PageModel = apps.get_model(app, model)
        convert_stream_field_rich_text_blocks(PageModel, field_name, convert_html_to_json)

    return migrations.RunPython(migrate_rich_text_blocks, reverse_migrate_rich_text_blocks)


def convert_stream_field_rich_text_blocks(model, field_name, convert_function):
    """
    go over each instance and convert field value
    :param model: Wagtail Page model class
    :param field_name: string
    :param convert_function: convert_json_to_html or convert_html_to_json expected
    :return: None
    """
    for instance in model.objects.all():
        value = getattr(instance, field_name)
        # conversion happens in place in nested structure so no need to set field value
        convert_rich_text_block_value(value, convert_function)
        instance.save(update_fields=[field_name])


def convert_rich_text_block_value(value, convert_function):
    """
    convert value for each found RichTextBlock
    :param value: StreamField value (instance of StreamValue)
    :param convert_function: convert_json_to_html or convert_html_to_json expected
    :return: None
    """
    # A way to traverse down to nested stream field structure
    if isinstance(value, StreamValue.StreamChild):
        convert_rich_text_block_value(value.value, convert_function)
    elif isinstance(value, (list, StreamValue)):
        for val in value:
            convert_rich_text_block_value(val, convert_function)
    elif isinstance(value, StructValue):
        for val in value.values():
            convert_rich_text_block_value(val, convert_function)
    elif isinstance(value, RichText):
        # fix RichTextBlock value
        value.source = convert_function(value.source)


def convert_rich_text_field_migration(app, model, field_name):
    """
    Creates a django RunPython migration to migrate json content of DraftailTextField to html content of RichTextField

    Usage:
    Create empty data migration and put line similar to this for each rich text field that needs to be migrated
    ```
        operations = [
            convert_rich_text_field_migration('core', 'BasePage', 'intro'),
        ]
    ```
    :param app: model app name as in django content type
    :param model: model name as in django content type
    :param field_name: rich text field name to migrate values for
    :return: django.db.migrations.operations.special.RunPython
    """
    def migrate_rich_text(apps, schema_editor):
        PageModel = apps.get_model(app, model)
        convert_model_field(PageModel, field_name, convert_json_to_html)

    def reverse_migrate_rich_text(apps, schema_editor):
        PageModel = apps.get_model(app, model)
        convert_model_field(PageModel, field_name, convert_html_to_json)

    return migrations.RunPython(migrate_rich_text, reverse_migrate_rich_text)


def convert_model_field(model, field_name, convert_function):
    """
    go over each instance and convert field value
    :param model: Wagtail Page model class
    :param field_name: string
    :param convert_function: convert_json_to_html or convert_html_to_json expected
    :return: None
    """
    for instance in model.objects.all():
        value = getattr(instance, field_name)
        value = convert_function(value)
        setattr(instance, field_name, value)
        instance.save(update_fields=[field_name])


def convert_json_to_html(value, features=None, options=None):
    """
    convert draftail API json to wagtail internal html content
    :param value: string of json from db
    :param features: optimal list of features as configured on RichTextField
    :param options: optimal dictionary of features options as configured on RichTextField
    :return: string of html content
    """
    converter = get_converter(features, options)
    return converter.to_database_format(value)


def convert_html_to_json(value, features=None, options=None):
    """
    convert wagtail internal html content to draftail API json
    :param value: string of html content
    :param features: optimal list of features as configured on RichTextField
    :param options: optimal dictionary of features options as configured on RichTextField
    :return: string of json from db
    """
    converter = get_converter(features, options)
    return converter.from_database_format(value)


def get_converter(features=None, options=None):
    """
    Initialize rich text converter in the same way as in wagtail.admin.rich_text.editors.draftail.DraftailRichTextArea
    :param features: optimal list of features as configured on RichTextField
    :param options: optimal dictionary of features options as configured on RichTextField
    :return: wagtail.admin.rich_text.converters.contentstate.ContentstateConverter
    """
    options = options or {}
    if features is None:
        features = feature_registry.get_default_features()
    for feature in features:
        plugin = feature_registry.get_editor_plugin('draftail', feature)
        if plugin:
            plugin.construct_options(options)

    converter = ContentstateConverter(features)
    return converter