md_to_gost/md2gost/renderable/image.py

import logging
from copy import copy
from io import BytesIO
from typing import Generator
from os import environ
import os.path

import requests
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.shared import Parented, Length, Mm
from docx.text.paragraph import Paragraph

from .caption import Caption, CaptionInfo
from .renderable import Renderable
from .requires_numbering import RequiresNumbering
from ..docx_svg import attach_svg_blip
from ..layout_tracker import LayoutState
from ..rendered_info import RenderedInfo
from ..sub_renderable import SubRenderable
from ..util import create_element

# Leave room under the figure so «Рисунок N — …» stays on the same page.
_CAPTION_RESERVE = Mm(12)


class Image(Renderable, RequiresNumbering):
    def __init__(
        self,
        parent: Parented,
        path: str,
        caption_info: CaptionInfo | None = None,
        svg_path: str | None = None,
        *,
        pixel_scale: float = 1.0,
    ):
        super().__init__("Рисунок")
        self._parent = parent
        self._caption_info = caption_info
        self._docx_paragraph = Paragraph(create_element("w:p"), parent)
        self._docx_paragraph.paragraph_format.space_before = 0
        self._docx_paragraph.paragraph_format.space_after = 0
        self._docx_paragraph.paragraph_format.first_line_indent = 0
        self._docx_paragraph.paragraph_format.line_spacing = 1
        self._docx_paragraph.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
        self._invalid = False
        self._native_width: Length | None = None
        self._native_height: Length | None = None

        run = self._docx_paragraph.add_run()

        if path.startswith("http"):
            bytesio = BytesIO()
            bytesio.write(requests.get(path).content)
            self._image = run.add_picture(bytesio)
        else:
            try:
                path = os.path.expanduser(path)
                if not os.path.isabs(path):
                    path = os.path.join(environ.get("WORKING_DIR", "."), path)
                self._image = run.add_picture(path)
            except FileNotFoundError:
                logging.warning(f"Invalid image path: {path}, skipping...")
                self._invalid = True
                self._image = None

        # High-res PlantUML/Kroki PNG: keep on-page size as if scale were 1.
        if not self._invalid and pixel_scale and pixel_scale > 1:
            self._image.width = Length(int(self._image.width / pixel_scale))
            self._image.height = Length(int(self._image.height / pixel_scale))

        if not self._invalid:
            self._native_width = Length(int(self._image.width))
            self._native_height = Length(int(self._image.height))

        if not self._invalid and svg_path:
            try:
                attach_svg_blip(run, self._image, svg_path)
            except Exception as exc:
                logging.warning("SVG blip attach failed (%s): %s", svg_path, exc)

        self._number = None
        self.landscape = bool(caption_info and caption_info.landscape)
        if caption_info and caption_info.unique_name:
            self.unique_name = caption_info.unique_name

        # Keep figure + caption together across Word pagination.
        if not self._invalid and caption_info is not None:
            self._docx_paragraph.paragraph_format.keep_with_next = True

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

    def _reset_native_size(self) -> None:
        if self._image is None or self._native_width is None or self._native_height is None:
            return
        self._image.width = Length(int(self._native_width))
        self._image.height = Length(int(self._native_height))

    def render(self, previous_rendered: RenderedInfo, layout_state: LayoutState)\
            -> Generator[RenderedInfo | SubRenderable, None, None]:
        if self._invalid:
            yield from []
            return

        # Re-fit from native size each time (Paragraph may have measured us earlier).
        self._reset_native_size()

        has_caption = self._caption_info is not None
        caption_reserve = _CAPTION_RESERVE if has_caption else Length(0)
        max_w = layout_state.max_width
        max_h = Length(max(0, int(layout_state.max_height) - int(caption_reserve)))

        # limit width
        if self._image.width > max_w:
            height_by_width = self._image.height / self._image.width
            self._image.width = max_w
            self._image.height = Length(self._image.width * height_by_width)

        # limit height (leave room for caption on the same page)
        if self._image.height > max_h:
            width_by_height = self._image.width / self._image.height
            self._image.height = max_h
            self._image.width = Length(self._image.height * width_by_height)

        height = self._image.height
        need = Length(int(height) + int(caption_reserve))
        remaining = layout_state.remaining_page_height
        if remaining < need:
            if self.landscape:
                # Already on a fresh landscape section — shrink instead of soft page-break
                # (soft break + section break = empty landscape page before the figure).
                avail = Length(max(0, int(remaining) - int(caption_reserve)))
                if avail > 0 and self._image.height > avail:
                    ratio = int(avail) / int(self._image.height)
                    self._image.height = avail
                    self._image.width = Length(int(self._image.width * ratio))
                height = self._image.height
            else:
                height = Length(int(height) + int(remaining))

        yield (rendered_image := RenderedInfo(self._docx_paragraph, Length(height)))

        layout_state.add_height(rendered_image.height)

        caption = Caption(self._parent, "Рисунок", self._caption_info, self._number, False)
        caption.center()

        yield from caption.render(rendered_image, copy(layout_state))