plugin.py

import os
import re
import shlex
import textwrap
from dataclasses import dataclass
from typing import List

import mkdocs.config.config_options
from mkdocs.plugins import BasePlugin

from codeinclude.languages import get_lang_class
from codeinclude.resolver import select

RE_START = r"""(?x)
    ^
    (?P<leading_space>\s*)
    <!--codeinclude-->
    (?P<ignored_trailing_space>\s*)
    $
"""

RE_END = r"""(?x)
    ^
    (?P<leading_space>\s*)
    <!--/codeinclude-->
    (?P<ignored_trailing_space>\s*)
    $
"""

RE_SNIPPET = r"""(?xm)
    ^
    (?P<leading_space>\s*)
    \[(?P<title>[^\]]*)\]\((?P<filename>[^)]+)\)
    ([\t\n ]+(?P<params>[\w:-]+))?
    (?P<ignored_trailing_space>\s*)
    $
"""


@dataclass
class CodeIncludeBlock(object):
    first_line_index: int
    last_line_index: int
    content: str


@dataclass
class Replacement(object):
    first_line_index: int
    last_line_index: int
    content: str


class CodeIncludePlugin(BasePlugin):

    config_scheme = (
        (
            "title_mode",
            mkdocs.config.config_options.Choice(
                choices=["none", "legacy_pymdownx.superfences", "pymdownx.tabbed", "mkdocs-material"],
                default="pymdownx.tabbed",
            ),
        ),
    )

    def on_page_markdown(self, markdown, page, config, site_navigation=None, **kwargs):
        """Provide a hook for defining functions from an external module"""

        blocks = self.find_code_include_blocks(markdown)
        substitutes = self.get_substitutes(blocks, page)
        return self.substitute(markdown, substitutes)

    def find_code_include_blocks(self, markdown: str) -> List[CodeIncludeBlock]:
        ci_blocks = list()
        first = -1
        in_block = False
        lines = markdown.splitlines()
        for index, line in enumerate(lines):
            if re.match(RE_START, lines[index]):
                if in_block:
                    raise ValueError(
                        f"Found two consecutive code-include starts: at lines {first} and {index}"
                    )
                first = index
                in_block = True
            elif re.match(RE_END, lines[index]):
                if not in_block:
                    raise ValueError(
                        f"Found code-include end without preceding start at line {index}"
                    )
                last = index
                content = "\n".join(lines[first : last + 1])
                ci_blocks.append(CodeIncludeBlock(first, last, content))
                in_block = False
        return ci_blocks

    def get_substitutes(
        self, blocks: List[CodeIncludeBlock], page
    ) -> List[Replacement]:
        replacements = list()
        for ci_block in blocks:
            replacement_content = ""
            for snippet_match in re.finditer(RE_SNIPPET, ci_block.content):
                title = snippet_match.group("title")
                filename = snippet_match.group("filename")
                indent = snippet_match.group("leading_space")
                raw_params = snippet_match.group("params")

                if raw_params:
                    params = dict(token.split(":") for token in shlex.split(raw_params))
                    lines = params.get("lines", "")
                    block = params.get("block", "")
                    inside_block = params.get("inside_block", "")
                else:
                    lines = ""
                    block = ""
                    inside_block = ""

                code_block = self.get_substitute(
                    page, title, filename, lines, block, inside_block
                )
                # re-indent
                code_block = re.sub("^", indent, code_block, flags=re.MULTILINE)

                replacement_content += code_block
            replacements.append(
                Replacement(
                    ci_block.first_line_index,
                    ci_block.last_line_index,
                    replacement_content,
                )
            )
        return replacements

    def get_substitute(self, page, title, filename, lines, block, inside_block):
        # Compute the fence header
        lang_code = get_lang_class(filename)
        header = lang_code
        title = title.strip()

        # Select the code content
        page_parent_dir = os.path.dirname(page.file.abs_src_path)
        import_path = os.path.join(page_parent_dir, filename)
        # Always use UTF-8, as it is the recommended default for source file encodings.
        with open(import_path, encoding="UTF-8") as f:
            content = f.read()

        selected_content = select(
            content, lines=lines, block=block, inside_block=inside_block
        )

        dedented = textwrap.dedent(selected_content)

        if self.config.get("title_mode") == "pymdownx.tabbed" and len(title) > 0:
            # Newest version of pymdownx requires everything to be indented 4-spaces.
            dedented = "".join([f'    {line}\n' for line in dedented.split('\n')])
            return f"""
=== "{title}"
    ```{header}
{dedented}
    ```

"""
        elif (
            self.config.get("title_mode") == "legacy_pymdownx.superfences"
            and len(title) > 0
        ):
            return f"""
```{header} tab="{title}"
{dedented}
```

"""
        elif (
            self.config.get("title_mode") == "mkdocs-material"
        ):
            return f"""
```{header} title="{title}"
{dedented}
```

"""     
        else:
            return f"""
```{header}
{dedented}
```

"""

    def substitute(self, markdown: str, substitutes: List[Replacement]) -> str:
        substitutes_by_first_line = dict()
        # Index substitutes by the first line
        for s in substitutes:
            substitutes_by_first_line[s.first_line_index] = s

        # Perform substitutions
        result = ""
        index = 0
        lines = markdown.splitlines()
        while index < len(lines):
            if index in substitutes_by_first_line.keys():
                # Replace the codeinclude fragment starting at this line
                substitute = substitutes_by_first_line[index]
                result += substitute.content
                index = substitute.last_line_index
            else:
                # Keep the input line
                result += lines[index] + "\n"
            index += 1
        return result