import markdown
from markdown.treeprocessors import Treeprocessor


class StyleTreeprocessor(Treeprocessor):

    def run_processor(self, node):
        for child in node:
            if child.tag == "table":
                child.set("class", "markdown-table")
            if child.tag == "th" or child.tag == "td":
                child.set("class", "markdown-table-cell")
            if child.tag == "img":
                child.set("class", "img-fluid")
            self.run_processor(child)
        return node

    def run(self, root):
        self.run_processor(root)
        return root


class MarkdownFormatExtension(markdown.Extension):
    # md_ globals deprecated - see here:
    def extendMarkdown(self, md):
        md.treeprocessors.register(
            StyleTreeprocessor(),
            'StyleTreeprocessor',
            10
        )