md_to_gost/md2gost/renderable/listing.py

from copy import copy
import os
from typing import Generator, Callable

from docx.shared import Length, Pt, RGBColor, Twips

from pygments import highlight
from pygments.formatter import Formatter
from pygments.lexers import get_lexer_by_name

from .caption import Caption, CaptionInfo
from .page_break import PageBreak
from .paragraph import Paragraph
from .renderable import Renderable
from .requires_numbering import RequiresNumbering
from ..docx_elements import create_table, create_table_row, create_table_cell, set_table_box_borders, _twips
from ..layout_tracker import LayoutState
from ..profiles import DEFAULT_LISTING_CONTINUATION, LISTING_CONTINUATION_MODES
from ..rendered_info import RenderedInfo
from ..sub_renderable import SubRenderable

_WORD_PAGED_MODES = frozenset({"off", "soft", "word"})


class DocxParagraphPygmentsFormatter(Formatter):
    def __init__(self, paragraphs: list[Paragraph], creator: Callable[[], Paragraph], **options):
        Formatter.__init__(self, style="sas", **options)
        self._creator = creator
        self._paragraphs = paragraphs
        self._styles = {}

        for token, style in self.style:
            self._styles[token] = style

    def _add_run_to_last_paragraph(self, text, style):
        self._paragraphs[-1].add_run(text, style["bold"] or None, style["italic"] in style or None,
                                     RGBColor.from_string(style['color']) if style['color'] else None)

    def format(self, tokensource, outfile):
        self._paragraphs.append(self._creator())
        for ttype, value in tokensource:
            style = self._styles[ttype]
            lines = iter(value.split("\n"))
            self._add_run_to_last_paragraph(next(lines), style)
            for line in lines:
                self._paragraphs.append(self._creator())
                self._add_run_to_last_paragraph(line, style)
        self._paragraphs.pop(-1)  # remove last empty line


LISTING_OFFSET = Pt(31) - Twips(108 * 2)  # todo: fix


class Listing(Renderable, RequiresNumbering):
    def __init__(self, parent, language: str, caption_info: CaptionInfo):
        super().__init__("Листинг")
        self._caption_info = caption_info
        self._language = language
        self._parent = parent
        self._continuation_mode = DEFAULT_LISTING_CONTINUATION
        self.paragraphs: list[Paragraph] = []
        self._number = None
        if caption_info and caption_info.unique_name:
            self.unique_name = caption_info.unique_name

    def _create_table(self, parent, width: Length):
        # Kept for tests / callers that still expect the helper; render uses multi-row.
        left_margin = Twips(int(
            parent.part.styles["Normal Table"]._element.xpath("w:tblPr/w:tblCellMar/w:left")[0].attrib[
                "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}w"]))
        right_margin = Twips(int(
            parent.part.styles["Normal Table"]._element.xpath("w:tblPr/w:tblCellMar/w:right")[0].attrib[
                "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}w"]))

        return create_table(
            parent, 1, 1,
            Twips(_twips(width) + _twips(left_margin) + _twips(right_margin)),
        )

    def set_text(self, text: str):
        def create_paragraph() -> Paragraph:
            paragraph = Paragraph(self._parent)
            paragraph.style = "Code"
            return paragraph

        text = text.removesuffix("\n")

        if self._language and "SYNTAX_HIGHLIGHTING" in os.environ and os.environ["SYNTAX_HIGHLIGHTING"] == "1":
            formatter = DocxParagraphPygmentsFormatter(self.paragraphs, lambda: create_paragraph())
            highlight(text, get_lexer_by_name(self._language), formatter)
        else:
            for line in text.removesuffix("\n").split("\n"):
                paragraph = create_paragraph()
                paragraph.add_run(line)
                self.paragraphs.append(paragraph)

    def set_number(self, number: str):
        self._number = number

    def set_continuation_mode(self, mode: str) -> None:
        if mode not in LISTING_CONTINUATION_MODES:
            raise ValueError(
                f"listing continuation must be one of {LISTING_CONTINUATION_MODES}, got {mode!r}"
            )
        self._continuation_mode = mode

    def _should_split_fragment(
        self, line_height, layout_state: LayoutState, lines_in_fragment: int
    ) -> bool:
        if self._continuation_mode in _WORD_PAGED_MODES:
            return False
        if lines_in_fragment == 0:
            return False
        return line_height > layout_state.remaining_page_height

    def _make_continuation_paragraph(self) -> Paragraph:
        continuation_paragraph = Paragraph(self._parent)
        continuation_paragraph.add_run(f"Продолжение Листинга {self._number}")
        continuation_paragraph.style = "Caption Listing"
        continuation_paragraph.first_line_indent = 0
        return continuation_paragraph

    def _emit_page_break_and_optional_caption(
        self, layout_state: LayoutState
    ) -> Generator[RenderedInfo, None, None]:
        page_break_info = next(PageBreak(self._parent).render(None, layout_state))
        layout_state.add_height(page_break_info.height)
        yield page_break_info

        continuation_paragraph = self._make_continuation_paragraph()
        continuation_paragraph._docx_paragraph.paragraph_format.keep_with_next = True
        continuation_paragraph._docx_paragraph.paragraph_format.page_break_before = False
        info = next(continuation_paragraph.render(None, copy(layout_state)))
        layout_state.add_height(info.height)
        yield info

    def render(self, previous_rendered: RenderedInfo, layout_state: LayoutState)\
            -> Generator[RenderedInfo | SubRenderable, None, None]:
        caption_rendered_infos = list(
            Caption(self._parent, "Листинг", self._caption_info, self._number, True)
            .render(previous_rendered, copy(layout_state))
        )
        layout_state.add_height(sum([info.height for info in caption_rendered_infos]))
        yield from caption_rendered_infos

        # One code line = one table row so Word can Split at real page breaks (mode word).
        # Outer box only — no grid lines between lines (looks like text in a frame).
        table = create_table(self._parent, 0, 1, self._listing_width(layout_state))
        set_table_box_borders(table)
        previous = None

        table_height = Pt(1)  # table borders, 4 eights of point for each border
        lines_in_fragment = 0
        col_w = self._listing_width(layout_state)

        # legacy/caption: if first line doesn't fit, burn the rest of the page
        if self.paragraphs and self._continuation_mode not in _WORD_PAGED_MODES:
            paragraph_layout_state = copy(layout_state)
            paragraph_layout_state.max_width -= LISTING_OFFSET
            first = next(self.paragraphs[0].render(previous, paragraph_layout_state))
            if first.height + table_height > layout_state.remaining_page_height:
                table_height += layout_state.remaining_page_height
                layout_state.add_height(layout_state.remaining_page_height)

        for paragraph in self.paragraphs:
            paragraph_layout_state = copy(layout_state)
            paragraph_layout_state.max_width -= LISTING_OFFSET
            paragraph_rendered_info = next(paragraph.render(previous, paragraph_layout_state))

            if self._should_split_fragment(
                paragraph_rendered_info.height, layout_state, lines_in_fragment
            ):
                yield RenderedInfo(table, table_height)
                yield from self._emit_page_break_and_optional_caption(layout_state)

                table_height = Pt(1)
                table = create_table(self._parent, 0, 1, col_w)
                set_table_box_borders(table)
                previous = None
                lines_in_fragment = 0

                paragraph_layout_state = copy(layout_state)
                paragraph_layout_state.max_width -= LISTING_OFFSET
                paragraph_rendered_info = next(paragraph.render(previous, paragraph_layout_state))

            row = create_table_row(table, header=False)
            cell = create_table_cell(row, col_w)
            cell._element.append(paragraph_rendered_info.docx_element._element)
            row._element.append(cell._element)
            table._element.append(row._element)

            layout_state.add_height(paragraph_rendered_info.height)
            table_height += paragraph_rendered_info.height
            lines_in_fragment += 1

            previous = paragraph_rendered_info

        yield RenderedInfo(table, table_height)

    def _listing_width(self, layout_state: LayoutState) -> Length:
        left_margin = Twips(int(
            self._parent.part.styles["Normal Table"]._element.xpath("w:tblPr/w:tblCellMar/w:left")[0].attrib[
                "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}w"]))
        right_margin = Twips(int(
            self._parent.part.styles["Normal Table"]._element.xpath("w:tblPr/w:tblCellMar/w:right")[0].attrib[
                "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}w"]))
        return Twips(_twips(layout_state.max_width) + _twips(left_margin) + _twips(right_margin))