"""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[^\s]+)\n(?P.*?)~~~', 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)