Facilitate extending HTMLRenderer to add attributes

mistletoe/html_renderer.py

@@ -5,6 +5,7 @@
 import re
 import sys
 from itertools import chain
+from typing import Dict, Union
 from urllib.parse import quote
 from mistletoe import block_token
 from mistletoe import span_token
@@ -16,12 +17,31 @@
 else:
     import html
 
+Attributes = Union[str, Dict[str, str]]
+"""
+Can be either a string or a dictionary mapping strings to strings.
+
+If it's a string, you need to take care of HTML escaping yourself.
+
+If it's a dictionary, the values are HTML escaped for you.
+"""
 
 class HTMLRenderer(BaseRenderer):
     """
     HTML renderer class.
 
     See mistletoe.base_renderer module for more info.
+
+    Some methods of the `HTMLRenderer` class take an optional `Attributes`
+    argument letting you easily add custom HTML attributes, for example:
+
+    .. code:: python
+
+        class MyRenderer(HTMLRenderer):
+            def render_heading(self, token):
+                return super().render_heading(
+                    token, attr={'id': self.render_to_plain(token).replace(' ', '-')}
+                )
     """
     def __init__(self, *extras):
         """
@@ -73,24 +93,39 @@ def render_image(self, token: span_token.Image) -> str:
             title = ''
         return template.format(token.src, self.render_to_plain(token), title)
 
-    def render_link(self, token: span_token.Link) -> str:
-        template = '<a href="{target}"{title}>{inner}</a>'
+    def render_link(self, token: span_token.Link, attr: Attributes = None) -> str:
+        """
+        Renders the given `Link`.
+
+        :Parameters:
+            attr
+                The ``href`` attribute is set regardless of what attributes you pass.
+        """
+        template = '<a href="{target}"{attr}>{inner}</a>'
         target = self.escape_url(token.target)
         if token.title:
-            title = ' title="{}"'.format(self.escape_html(token.title))
-        else:
-            title = ''
+            if attr is None:
+                attr = dict(title=token.title)
+            elif isinstance(attr, dict):
+                attr['title'] = token.title
         inner = self.render_inner(token)
-        return template.format(target=target, title=title, inner=inner)
+        return template.format(target=target, inner=inner, attr=self._render_attributes(attr))
 
-    def render_auto_link(self, token: span_token.AutoLink) -> str:
-        template = '<a href="{target}">{inner}</a>'
+    def render_auto_link(self, token: span_token.AutoLink, attr: Attributes = None) -> str:
+        """
+        Renders the given `AutoLink`.
+
+        :Parameters:
+            attr
+                The ``href`` attribute is set regardless of what attributes you pass.
+        """
+        template = '<a href="{target}"{attr}>{inner}</a>'
         if token.mailto:
             target = 'mailto:{}'.format(token.target)
         else:
             target = self.escape_url(token.target)
         inner = self.render_inner(token)
-        return template.format(target=target, inner=inner)
+        return template.format(target=target, inner=inner, attr=self._render_attributes(attr))
 
     def render_escape_sequence(self, token: span_token.EscapeSequence) -> str:
         return self.render_inner(token)
@@ -102,13 +137,13 @@ def render_raw_text(self, token: span_token.RawText) -> str:
     def render_html_span(token: span_token.HTMLSpan) -> str:
         return token.content
 
-    def render_heading(self, token: block_token.Heading) -> str:
-        template = '<h{level}>{inner}</h{level}>'
+    def render_heading(self, token: block_token.Heading, attr: Attributes = None) -> str:
+        template = '<h{level}{attr}>{inner}</h{level}>'
         inner = self.render_inner(token)
-        return template.format(level=token.level, inner=inner)
+        return template.format(level=token.level, inner=inner, attr=self._render_attributes(attr))
 
-    def render_quote(self, token: block_token.Quote) -> str:
-        elements = ['<blockquote>']
+    def render_quote(self, token: block_token.Quote, attr: Attributes = None) -> str:
+        elements = ['<blockquote{}>'.format(self._render_attributes(attr))]
         self._suppress_ptag_stack.append(False)
         elements.extend([self.render(child) for child in token.children])
         self._suppress_ptag_stack.pop()
@@ -120,27 +155,39 @@ def render_paragraph(self, token: block_token.Paragraph) -> str:
             return '{}'.format(self.render_inner(token))
         return '<p>{}</p>'.format(self.render_inner(token))
 
-    def render_block_code(self, token: block_token.BlockCode) -> str:
+    def render_block_code(self, token: block_token.BlockCode, attr: Attributes = None) -> str:
+        """
+        Renders the given `BlockCode`.
+
+        :Parameters:
+            attr
+                Defaults to ``class="language-{token.language}"`` for code blocks that have a language.
+        """
         template = '<pre><code{attr}>{inner}</code></pre>'
-        if token.language:
-            attr = ' class="{}"'.format('language-{}'.format(self.escape_html(token.language)))
-        else:
-            attr = ''
+        if token.language and attr is None:
+            attr = 'class="{}"'.format('language-{}'.format(self.escape_html(token.language)))
         inner = html.escape(token.children[0].content)
-        return template.format(attr=attr, inner=inner)
+        return template.format(inner=inner, attr=self._render_attributes(attr))
 
-    def render_list(self, token: block_token.List) -> str:
-        template = '<{tag}{attr}>\n{inner}\n</{tag}>'
+    def render_list(self, token: block_token.List, attr: Attributes = None) -> str:
+        """
+        Renders the given `List`.
+
+        :Parameters:
+            attr
+                The ``start`` attribute is set regardless of what attributes you pass.
+        """
+        template = '<{tag}{start}{attr}>\n{inner}\n</{tag}>'
         if token.start is not None:
             tag = 'ol'
-            attr = ' start="{}"'.format(token.start) if token.start != 1 else ''
+            start = ' start="{}"'.format(token.start) if token.start != 1 else ''
         else:
             tag = 'ul'
-            attr = ''
+            start = ''
         self._suppress_ptag_stack.append(not token.loose)
         inner = '\n'.join([self.render(child) for child in token.children])
         self._suppress_ptag_stack.pop()
-        return template.format(tag=tag, attr=attr, inner=inner)
+        return template.format(tag=tag, start=start, inner=inner, attr=self._render_attributes(attr))
 
     def render_list_item(self, token: block_token.ListItem) -> str:
         if len(token.children) == 0:
@@ -154,12 +201,12 @@ def render_list_item(self, token: block_token.ListItem) -> str:
                 inner_template = inner_template[:-1]
         return '<li>{}</li>'.format(inner_template.format(inner))
 
-    def render_table(self, token: block_token.Table) -> str:
+    def render_table(self, token: block_token.Table, attr: Attributes = None) -> str:
         # This is actually gross and I wonder if there's a better way to do it.
         #
         # The primary difficulty seems to be passing down alignment options to
         # reach individual cells.
-        template = '<table>\n{inner}</table>'
+        template = '<table{attr}>\n{inner}</table>'
         if hasattr(token, 'header'):
             head_template = '<thead>\n{inner}</thead>\n'
             head_inner = self.render_table_row(token.header, is_header=True)
@@ -168,7 +215,7 @@ def render_table(self, token: block_token.Table) -> str:
         body_template = '<tbody>\n{inner}</tbody>\n'
         body_inner = self.render_inner(token)
         body_rendered = body_template.format(inner=body_inner)
-        return template.format(inner=head_rendered+body_rendered)
+        return template.format(inner=head_rendered+body_rendered, attr=self._render_attributes(attr))
 
     def render_table_row(self, token: block_token.TableRow, is_header=False) -> str:
         template = '<tr>\n{inner}</tr>\n'
@@ -216,3 +263,12 @@ def escape_url(raw: str) -> str:
         Escape urls to prevent code injection craziness. (Hopefully.)
         """
         return html.escape(quote(html.unescape(raw), safe='/#:()*?=%@+,&;'))
+
+    @classmethod
+    def _render_attributes(cls, attrs: Attributes = None) -> str:
+        if attrs is None:
+            return ''
+        if isinstance(attrs, str):
+            return ' ' + attrs
+        else:
+            return ' ' + ' '.join('{}="{}"'.format(k, cls.escape_html(v)) for k, v in attrs.items())

test/test_html_renderer.py

@@ -153,3 +153,30 @@ def test_footnote_link(self):
         token = Document(['[name][foo]\n', '\n', '[foo]: target\n'])
         output = '<p><a href="target">name</a></p>\n' 
         self.assertEqual(self.renderer.render(token), output)
+
+class CustomRenderer(HTMLRenderer):
+    def render_heading(self, token):
+        return super().render_heading(
+            token, attr={'id': self.render_to_plain(token).replace(' ', '-'), 'class': 'fancy-heading'}
+        )
+
+    def render_list(self, token):
+        return super().render_list(token, 'class=fancy-list')
+
+
+class TestCustomAttributes(TestCase):
+    def test_heading_ids(self):
+        from mistletoe import Document
+        with CustomRenderer() as renderer:
+            doc = Document('# hello "world"')
+            # test that attribute values are HTML-escaped
+            output = '<h1 id="hello-&quot;world&quot;" class="fancy-heading">hello &quot;world&quot;</h1>\n'
+            self.assertEqual(renderer.render(doc), output)
+
+    def test_list(self):
+        from mistletoe import Document
+        with CustomRenderer() as renderer:
+            doc = Document('2. foo')
+            # test that start attribute is preserved
+            output = '<ol start="2" class=fancy-list>\n<li>foo</li>\n</ol>\n'
+            self.assertEqual(renderer.render(doc), output)