Untitled

from django.utils.html import format_html
from wagtail.core import hooks
from wagtail.core.rich_text import features
from wagtail.core.rich_text.contentstate import ContentState
from wagtail.core.rich_text.contentstate import ContentStateToHTMLConverter
from wagtail.core.rich_text.features import InlineStyleFeature

ALIGN_LEFT = 'text-align-left'
ALIGN_CENTER = 'text-align-center'
ALIGN_RIGHT = 'text-align-right'

class AlignmentHandler:
    """
    Handler for converting text alignment from HTML to database format and vice versa.
    """
    def from_database_format(self, value):
        """
        Converts text alignment from database format to HTML.
        """
        return format_html('<div class="{}">{}</div>', value, '{}')

    def to_database_format(self, value):
        """
        Converts text alignment from HTML to database format.
        """
        if value and value in [ALIGN_LEFT, ALIGN_CENTER, ALIGN_RIGHT]:
            return value
        return None

class AlignmentFeature(InlineStyleFeature):
    """
    Feature for adding text alignment to the RichTextField.
    """
    def __init__(self):
        super().__init__(
            'alignment',
            label='Alignment',
            icon='fa-align-left',
            style_choices=[
                ('', 'Normal'),
                (ALIGN_LEFT, 'Left'),
                (ALIGN_CENTER, 'Center'),
                (ALIGN_RIGHT, 'Right'),
            ]
        )

    def get_css_classes(self, style):
        """
        Returns the CSS classes to apply based on the selected style.
        """
        if style in [ALIGN_LEFT, ALIGN_CENTER, ALIGN_RIGHT]:
            return 'align ' + style
        return ''

    def get_element(self, style, value):
        """
        Returns the HTML element to use based on the selected style.
        """
        class_name = self.get_css_classes(style)
        return format_html('<div class="{}">{}</div>', class_name, self.icon)

# Register the plugin with Wagtail
@hooks.register('register_rich_text_features')
def register_alignment_feature(features):
    """
    Registers the alignment feature with Wagtail.
    """
    feature_name = 'alignment'
    type_ = 'text-align'
    tag = 'div'

    # Define the converter rules
    converter_rules = {
        'from_database_format': {
            'div[class]': AlignmentHandler(),
        },
        'to_database_format': {
            'style_map': {
                'text-align-left': ALIGN_LEFT,
                'text-align-center': ALIGN_CENTER,
                'text-align-right': ALIGN_RIGHT
            }
        },
        'to_html': {
            'block_map': {
                type_: ContentStateToHTMLConverter.ElementHandler(tag)
            }
        }
    }

    # Register the feature
    feature = AlignmentFeature()
    features.register_editor_plugin(feature_name, feature)
    features.register_converter_rule('contentstate', feature_name, converter_rules)
    features.default_features.append(feature_name)


# Example usage in a model
from wagtail.core.models import Page
from wagtail.core.fields import RichTextField

class MyPage(Page):
    body = RichTextField(features=['bold', 'italic', 'link', 'alignment'])