richtext filter for Wagtail that creates relative links

rich_text.py

"""
This overrides some of the default behaviour of richtext from Wagtail
"""

from wagtail.wagtailcore.rich_text import get_link_handler, extract_attrs, get_embed_handler, LINK_HANDLERS, FIND_A_TAG, FIND_EMBED_TAG
from wagtail.wagtailcore.rich_text import PageLinkHandler as OriginalPageLinkHandler
from wagtail.wagtailcore.models import Page

from django.utils.html import escape


class PageLinkHandler(OriginalPageLinkHandler):
    """
    Similar to Wagtail's PageLinkHandler but can build relative URLs if passed current_page
    """

    @staticmethod
    def expand_db_attributes(attrs, for_editor, current_page):
        try:
            page = Page.objects.get(id=attrs['id'])

            if for_editor:
                editor_attrs = 'data-linktype="page" data-id="%d" ' % page.id
                parent_page = page.get_parent()
                if parent_page:
                    editor_attrs += 'data-parent-id="%d" ' % parent_page.id
            else:
                editor_attrs = ''

            current_site = current_page.get_site() if current_page else None
            if current_page and current_site == page.specific.get_site():
                url = page.specific.relative_url(current_site)
            else:
                url = page.specific.url

            return '<a %shref="%s">' % (editor_attrs, escape(url))
        except Page.DoesNotExist:
            return "<a>"
        except Exception as e:
            import IPython; IPython.embed(using=False)


# Override default PageLinkHandler with our own
LINK_HANDLERS['page'] = PageLinkHandler


def expand_db_html(html, for_editor=False, current_page=None):
    """
    Expand database-representation HTML into proper HTML usable in either
    templates or the rich text editor
    """
    def replace_a_tag(m):
        attrs = extract_attrs(m.group(1))
        if 'linktype' not in attrs:
            # return unchanged
            return m.group(0)
        handler = get_link_handler(attrs['linktype'])
        if handler == PageLinkHandler:
            return handler.expand_db_attributes(attrs, for_editor, current_page)
        return handler.expand_db_attributes(attrs, for_editor)

    def replace_embed_tag(m):
        attrs = extract_attrs(m.group(1))
        handler = get_embed_handler(attrs['embedtype'])
        if isinstance(handler, PageLinkHandler):
            return handler.expand_db_attributes(attrs, for_editor, current_page)
        return handler.expand_db_attributes(attrs, for_editor)

    html = FIND_A_TAG.sub(replace_a_tag, html)
    html = FIND_EMBED_TAG.sub(replace_embed_tag, html)
    return html

templatetags.py

from django import template
from wagtail.wagtailcore.rich_text import RichText

from .rich_text import expand_db_html

register = template.Library()


@register.filter
def richtext(value, current_page=None):
    """
    Similar to Wagtail's richtext but we can also pass the current page to create relative links instead
    of absolute when possible
    """
    if current_page and not isinstance(current_page, Page):
        raise Exception("Filter's first argument must be the current page - aka the page used in the template")

    if isinstance(value, RichText):
        # passing a RichText value through the |richtext filter should have no effect
        return value
    elif value is None:
        html = ''
    else:
        html = expand_db_html(value, current_page=current_page)

    return mark_safe('<div class="rich-text">' + html + '</div>')