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.entities import EntityElementHandler
from wagtail.core.rich_text import Expansion
from wagtail.core.rich_text import LinkHandler
from wagtail.core.rich_text import RichText
from wagtail.core.rich_text import Source

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

@hooks.register('register_rich_text_features')
def register_alignment_feature(features):
    """
    Register the alignment feature for the RichTextField.
    """
    feature_name = 'alignment'
    type_ = 'ALIGNMENT'
    tag = 'div'
    handler = AlignmentHandler

    # Register the alignment feature
    feature = features.register_editor_plugin(
        'draftail', feature_name, features.InlineEntityFeature(type_, tag, handler)
    )
    # Add the feature to the toolbar
    features.default_features.append(feature)
    # Add the feature to the feature panels
    features.register_feature(features.Feature(type_, feature_name, label='Alignment'))

class AlignmentHandler(EntityElementHandler):
    """
    Handler for converting text alignment from HTML to database format and vice versa.
    """
    mutability = 'IMMUTABLE'

    def get_attribute_data(self, attrs):
        """
        Returns the text alignment value from the HTML attributes.
        """
        return attrs.get('data-alignment')

    def get_tag(self, attrs):
        """
        Returns the HTML tag to use for this entity.
        """
        return 'span'

    def expand_entity(self, instance, start_pos, end_pos):
        """
        Returns the HTML representation of the text alignment entity.
        """
        alignment = instance['data']
        return format_html('<div class="{}">{}</div>', alignment, instance['content'])

@hooks.register('construct_whitelister_element_rules')
def whitelister_element_rules():
    """
    Add the alignment tag to the whitelisted elements.
    """
    return {ALIGN_LEFT: {'attributes': ['data-alignment'], 'check': check_alignment},
            ALIGN_CENTER: {'attributes': ['data-alignment'], 'check': check_alignment},
            ALIGN_RIGHT: {'attributes': ['data-alignment'], 'check': check_alignment}}

def check_alignment(attrs):
    """
    Check the attributes for the alignment tag.
    """
    return 'text-align-' in attrs.get('class', '')

class AlignmentFeature(features.Feature):
    """
    Feature for adding text alignment to the RichTextField.
    """
    def __init__(self):
        super().__init__(
            'alignment',
            label='Alignment',
            icon='fa-align-left',
            type_='ALIGNMENT',
            default_alignment='',
            selector='div',
            style=dict.fromkeys([ALIGN_LEFT, ALIGN_CENTER, ALIGN_RIGHT], {}),
            format='div'
        )

    def clean(self, value):
        """
        Cleans the text alignment value.
        """
        alignment = value.get('alignment')
        if alignment and alignment in [ALIGN_LEFT, ALIGN_CENTER, ALIGN_RIGHT]:
            return {'alignment': alignment}
        return {}

    def render_panel(self):
        """
        Returns the HTML for the alignment panel.
        """
        return ''

@hooks.register('register_rich_text_features')
def register_custom_richtext_features(features):
    """
    Register custom richtext features.
    """
    features.register_feature(AlignmentFeature())