incorporeal-cms/incorporealcms/mdx/pydot.py

52 lines
1.7 KiB
Python

"""Serve dot diagrams inline."""
import base64
import logging
import re
import markdown
import pydot
logger = logging.getLogger(__name__)
class InlinePydot(markdown.Extension):
"""Wrap the markdown prepcoressor."""
def extendMarkdown(self, md):
"""Add InlinePydotPreprocessor to the Markdown instance."""
md.preprocessors.register(InlinePydotPreprocessor(md), 'dot_block', 100)
class InlinePydotPreprocessor(markdown.preprocessors.Preprocessor):
"""Identify dot codeblocks and run them through pydot."""
BLOCK_RE = re.compile(r'~~~pydot:(?P<filename>[^\s]+)\n(?P<content>.*?)~~~', re.DOTALL)
def run(self, lines):
"""Match and generate diagrams from dot code blocks."""
text = '\n'.join(lines)
for match in self.BLOCK_RE.finditer(text):
filename = match.group(1)
dot_string = match.group(2)
# use pydot to turn the text into pydot
graphs = pydot.graph_from_dot_data(dot_string)
if not graphs:
logger.debug("some kind of issue with parsed 'dot' %s", dot_string)
raise ValueError("error parsing dot text!")
# encode the image and provide as an inline image in markdown
encoded_image = base64.b64encode(graphs[0].create_png()).decode('ascii')
data_path = f'data:image/png;base64,{encoded_image}'
inline_image = f'![{filename}]({data_path})'
# replace the image in the output markdown
text = f'{text[:match.start()]}\n{inline_image}\n{text[match.end():]}'
return text.split('\n')
def makeExtension(*args, **kwargs):
"""Provide the extension to the markdown extension loader."""
return InlinePydot(*args, **kwargs)