wagtail_streamfield_pagerevision_migration.py

# -*- coding: utf-8 -*-
from __future__ import unicode_literals

import json

from django.core.serializers.json import DjangoJSONEncoder
from django.db import migrations, models

from wagtail.wagtailcore.rich_text import RichText

target_models = [
    'StandardPage',
    'StandardFormPage',
    'EventPage',
    'ResourcePage',
    'NewsItemPage',
]


def page_to_streamfield(page):
    changed = False
    for attr in ('body_en', 'body_fr', 'body_es'):
        field = getattr(page, attr)
        if field.raw_text and not field:
            setattr(page, attr, [('rich_text', {'rich_text': RichText(field.raw_text)})])
            changed = True
    return page, changed


def pagerevision_to_streamfield(revision_data):
    changed = False
    for attr in ('body_en', 'body_fr', 'body_es'):
        fieldvalue = revision_data.get(attr)
        if fieldvalue:
            try:
                json.loads(fieldvalue)
            except ValueError:
                revision_data[attr] = json.dumps(
                    [{
                        "value": {"rich_text": fieldvalue},
                        "type": "rich_text"
                    }],
                    cls=DjangoJSONEncoder)
                changed = True
    return revision_data, changed


def page_to_richtext(page):
    changed = False
    for attr in ('body_en', 'body_fr', 'body_es'):
        field = getattr(page, attr)
        if field.raw_text is None:
            raw_text = ''.join([
                child.value['rich_text'].source for child in field
                if child.block_type == 'rich_text'
            ])
            setattr(page, attr, raw_text)
            changed = True
    return page, changed


def pagerevision_to_richtext(revision_data):
    changed = False
    for attr in ('body_en', 'body_fr', 'body_es'):
        fieldvalue = revision_data.get(attr, 'definitely non-JSON string')
        if fieldvalue:
            try:
                field_data = json.loads(fieldvalue)
            except ValueError:
                # It's not apparently a streamfield. Leave it.
                pass
            else:
                raw_text = ''.join([
                    child['value']['rich_text'] for child in field_data
                    if child['type'] == 'rich_text'
                ])
                revision_data[attr] = raw_text
                changed = True
    return revision_data, changed


def convert(apps, schema_editor, page_converter, pagerevision_converter):
    for page_model in target_models:
        model = apps.get_model("home", page_model)
        for page in model.objects.all():

            page, changed = page_converter(page)
            if changed:
                page.save()

            for revision in page.revisions.all():
                revision_data = json.loads(revision.content_json)
                revision_data, changed = pagerevision_converter(revision_data)
                if changed:
                    revision.content_json = json.dumps(revision_data, cls=DjangoJSONEncoder)
                    revision.save()


def convert_to_streamfield(apps, schema_editor):
    return convert(apps, schema_editor, page_to_streamfield, pagerevision_to_streamfield)


def convert_to_richtext(apps, schema_editor):
    return convert(apps, schema_editor, page_to_richtext, pagerevision_to_richtext)


class Migration(migrations.Migration):

    dependencies = [
        ('home', 'insert_dependency_here'),
    ]

    operations = [
        migrations.RunPython(
            convert_to_streamfield,
            convert_to_richtext,
        ),
    ]