from __future__ import annotations import re from typing import TYPE_CHECKING, Any, Literal, cast from markdown import Extension from markdown.preprocessors import Preprocessor from auto_pytabs.core import ( Cache, VersionedCode, VersionTuple, get_version_requirements, version_code, ) if TYPE_CHECKING: from markdown import Markdown RGX_BLOCK_TOKENS = re.compile(r"(.*```py[\w\W]*)|(.*```)") RGX_PYTABS_DIRECTIVE = re.compile(r"") PyTabDirective = Literal["disable", "enable", "disable-block"] PYTAB_DIRECTIVES: set[PyTabDirective] = {"disable", "enable", "disable-block"} def _strip_indentation(lines: list[str]) -> tuple[list[str], str]: if not lines: return [], "" first_line = lines[0] indent_char = "" if first_line[0] in [" ", "\t"]: indent_char = first_line[0] indent = indent_char * (len(first_line) - len(first_line.lstrip(indent_char))) if indent: return [line.split(indent, 1)[1] if line else "" for line in lines], indent return lines, indent def _add_indentation(source: str | list[str], indentation: str) -> str: lines = source.splitlines() if isinstance(source, str) else source return "\n".join(indentation + line for line in lines) def _get_pytabs_directive(line: str) -> PyTabDirective | None: match = RGX_PYTABS_DIRECTIVE.match(line) if match: matched_directive = match.group(1).strip() if matched_directive in PYTAB_DIRECTIVES: return cast(PyTabDirective, matched_directive) raise RuntimeError(f"Invalid AutoPytabs directive: {matched_directive!r}") return None def _extract_code_blocks(lines: list[str]) -> tuple[list[str], dict[int, list[str]]]: in_block = False enabled = True new_lines: list[str] = [] to_transform = {} start = 0 for i, line in enumerate(lines): is_comment_line = False directive = _get_pytabs_directive(line) if directive: is_comment_line = True if directive == "disable": enabled = False elif directive == "enable": enabled = True match = RGX_BLOCK_TOKENS.match(line) if match: if match.group(1): in_block = True start = i elif match.group(2) and in_block: in_block = False block = lines[start : i + 1] block_directive = _get_pytabs_directive(lines[start - 1]) if enabled and block_directive != "disable-block": to_transform[len(new_lines)] = block new_lines.append("") else: new_lines.extend(block) else: new_lines.append(line) elif not in_block and not is_comment_line: new_lines.append(line) return new_lines, to_transform def _build_tab(title: str, body: list[str], selected: bool) -> str: out = f'==={"+" if selected else ""} "{title}"\n' out += _add_indentation(body, indentation=" ") out += "\n" return out def _build_tabs( *, versioned_code: VersionedCode, head: str, tail: str, tab_title_template: str, default_tab_version: VersionTuple, reverse_order: bool, ) -> str: out = [] for version, code in versioned_code.items(): version_string = f"{version[0]}.{version[1]}" lines = [head, *code.splitlines(), tail] tab_title = tab_title_template.format(min_version=version_string) out.append( _build_tab( title=tab_title, body=lines, selected=version == default_tab_version ) ) if reverse_order: out = reversed(out) # type: ignore[assignment] return "\n".join(out) def _convert_block( *, block: list[str], versions: list[VersionTuple], tab_title_template: str, cache: Cache | None, default_tab_strategy: Literal["highest", "lowest"], reverse_order: bool, ) -> str: block, indentation = _strip_indentation(block) head, *code_lines, tail = block code = "\n".join(code_lines) versioned_code = version_code(code, versions, cache=cache) if len(versioned_code) > 1: versions = list(versioned_code.keys()) default_tab_version = versions[-1 if default_tab_strategy == "highest" else 0] code = _build_tabs( versioned_code=versioned_code, head=head, tail=tail, tab_title_template=tab_title_template, default_tab_version=default_tab_version, reverse_order=reverse_order, ) else: code = "\n".join([head, versioned_code[versions[0]], tail]) code = _add_indentation(code, indentation) return code class UpgradePreprocessor(Preprocessor): def __init__( self, *args: Any, min_version: str, max_version: str, tab_title_template: str | None = None, cache: Cache | None = None, default_tab_strategy: Literal["highest", "lowest"] = "highest", reverse_order: bool = False, **kwargs: Any, ) -> None: self.min_version = VersionTuple.from_string(min_version) self.max_version = VersionTuple.from_string(max_version) self.versions = get_version_requirements(self.min_version, self.max_version) self.tab_title_template = tab_title_template or "Python {min_version}+" self.cache = cache self.default_tab_strategy = default_tab_strategy self.reverse_order = reverse_order super().__init__(*args, **kwargs) def run(self, lines: list[str]) -> list[str]: new_lines, to_transform = _extract_code_blocks(lines) output_lines = [] for i, line in enumerate(new_lines): block_to_transform = to_transform.get(i) if block_to_transform: transformed_block = _convert_block( block=block_to_transform, versions=self.versions, tab_title_template=self.tab_title_template, cache=self.cache, default_tab_strategy=self.default_tab_strategy, reverse_order=self.reverse_order, ).splitlines() output_lines.extend(transformed_block) else: output_lines.append(line) return output_lines class AutoPyTabsExtension(Extension): def __init__(self, *args: Any, cache: Cache | None, **kwargs: Any): self.config = { "min_version": ["3.7", "minimum version"], "max_version": ["3.11", "maximum version"], "tab_title_template": ["", "tab title format-string"], "default_tab": ["highest", "version tab to preselect"], "reverse_order": [False, "reverse the order of tabs"], } self.cache = cache super().__init__(*args, **kwargs) def extendMarkdown(self, md: Markdown) -> None: """Register the extension.""" self.md = md md.registerExtension(self) config = self.getConfigs() md.preprocessors.register( UpgradePreprocessor( min_version=config["min_version"], max_version=config["max_version"], tab_title_template=config["tab_title_template"], cache=self.cache, default_tab_strategy=config["default_tab"], reverse_order=config["reverse_order"], ), "auto_pytabs", 32, ) def makeExtension(**kwargs: Any) -> AutoPyTabsExtension: return AutoPyTabsExtension(**kwargs)