formula_img2latex/formula_img2latex/sources/base.py

"""Shared Markdown scanning and the Source interface."""

from __future__ import annotations

import re
from dataclasses import dataclass
from pathlib import Path
from urllib.parse import unquote

IMAGE_RE = re.compile(r"!\[((?:[^\]\\]|\\.)*)\]\(\s*<?([^)\s>]+)>?(?:\s+\"[^\"]*\")?\s*\)")
_FENCE_RE = re.compile(r"^\s*(```+|~~~+)")
_RASTER_EXT = (".png", ".gif", ".jpg", ".jpeg", ".svg", ".svgz", ".webp", ".bmp", ".tif", ".tiff")


@dataclass
class ImageRef:
    file: Path          # Markdown file (absolute)
    line: int           # 1-based line number
    start: int          # column span of the reference within the line
    end: int
    md: str             # the exact reference text
    alt: str
    target: str         # link target as written
    display: bool       # alone in its paragraph -> display math
    context: str        # nearby text, for prompts and review

    def asset_path(self) -> Path:
        return (self.file.parent / unquote(self.target)).resolve()


def scan_markdown(md_file: Path) -> list[ImageRef]:
    """All local raster/SVG image references outside code blocks."""
    lines = md_file.read_text(encoding="utf-8").split("\n")
    refs: list[ImageRef] = []
    fence: str | None = None
    for i, line in enumerate(lines):
        m = _FENCE_RE.match(line)
        if m:
            mark = m.group(1)[0] * 3
            if fence is None:
                fence = mark
            elif mark == fence:
                fence = None
            continue
        if fence is not None or "![" not in line:
            continue
        for im in IMAGE_RE.finditer(line):
            target = im.group(2)
            if "://" in target or target.startswith("data:") or not target.lower().endswith(_RASTER_EXT):
                continue
            rest = (line[: im.start()] + line[im.end():]).strip()
            prev_blank = i == 0 or not lines[i - 1].strip()
            next_blank = i + 1 >= len(lines) or not lines[i + 1].strip()
            display = rest in ("", ">") and prev_blank and next_blank and not line.lstrip().startswith("|")
            ctx = rest if rest and rest != ">" else (lines[i - 2].strip() if i >= 2 else "")
            refs.append(ImageRef(md_file, i + 1, im.start(), im.end(), im.group(0),
                                 im.group(1), target, display, ctx[:300]))
    return refs


class Source:
    """A converter output directory or file. Subclasses fill in the specifics."""

    kind = "?"

    def __init__(self, path: Path) -> None:
        self.path = path.resolve()

    @classmethod
    def detect(cls, path: Path) -> bool:
        raise NotImplementedError

    @property
    def root(self) -> Path:
        """Directory all relative paths in the workspace refer to."""
        return self.path if self.path.is_dir() else self.path.parent

    def markdown_files(self) -> list[Path]:
        raise NotImplementedError

    def describe(self) -> dict:
        return {"kind": self.kind, "path": str(self.path)}

    def image_refs(self) -> list[ImageRef]:
        refs: list[ImageRef] = []
        for md in self.markdown_files():
            refs.extend(scan_markdown(md))
        return refs

    def known_formula(self, ref: ImageRef) -> bool:
        return False

    def source_name(self, ref: ImageRef) -> str:
        return ref.target

    def original_bytes(self, ref: ImageRef) -> bytes:
        """Best available original of the formula (pdf2md re-crops from the vector PDF)."""
        return ref.asset_path().read_bytes()

    def font_hint(self) -> str | None:
        return None

    def math_delimiters(self) -> tuple[str, str, str, str]:
        return "$", "$", "$$", "$$"

    def language(self) -> str:
        return ""

    def format_math(self, tex: str, display: bool) -> str:
        il, ir, dl, dr = self.math_delimiters()
        body = tex.strip()
        if self.language().lower().startswith("de"):
            # German decimal comma: 3,6 -> 3{,}6, otherwise LaTeX sets "3, 6"
            body = re.sub(r"(?<=\d),(?=\d)", "{,}", body)
        if display:
            return f"{dl}{body}{dr}"
        return f"{il}{body}{ir}"