omnisvg-lite/incremental_svg.py

"""Инкрементальный SVG-парсер.

Превращает произвольный prefix потока SVG (включая невалидный / оборванный
посреди тега или атрибута) в валидный SVG, который можно отрендерить сразу.

Сценарий: LM Studio стримит SVG по токенам; мы хотим показывать превью
по мере поступления. В буфере в любой момент времени может быть:
  - незакрытый открывающий тег:       `<svg viewBox="0 0 64 64"`
  - оборванный атрибут:                `<path d="M12 32 C `
  - незакрытая кавычка:                `<rect fill="`
  - частично написанный тег:           `<path d="M12 32 C 40 20,`
  - не закрытый `<svg>`:               `<svg ...><circle cx="50"`

`parse_to_valid(prefix)`:
  1. Находит первый `<svg` в потоке (всё до него — мусор/reasoning/fence).
  2. Обрезает по первому `</svg>` (остальное — мусор/повторный SVG).
  3. Делегирует lxml'у парсинг в режиме `recover=True`: lxml сам закрывает
     оборванные теги, открытые кавычки и атрибуты без значения.
  4. Добавляет `xmlns="http://www.w3.org/2000/svg"` в корневой `<svg>`,
     если его нет.
  5. Гарантирует, что результат парсится `lxml.etree.fromstring()` без ошибок.

Использование:

    from incremental_svg import parse_to_valid, stream_to_valid

    # Один снапшот
    svg = parse_to_valid(raw_prefix)

    # Стрим: получаем валидный SVG после каждого нового чанка
    for snapshot in stream_to_valid(token_chunks):
        render_to_png(snapshot)
"""
from __future__ import annotations

import re
from typing import Iterable, Iterator

from lxml import etree


# ---------------------------------------------------------------------------
# Константы
# ---------------------------------------------------------------------------

SVG_NS = "http://www.w3.org/2000/svg"

# Дефолтный пустой SVG (xmlns обязателен — без него некоторые рендереры падают).
DEFAULT_EMPTY_SVG = f'<svg xmlns="{SVG_NS}"></svg>'

# Находим первое вхождение `<svg` (case-insensitive). Используется только
# позиция совпадения; сам тег обрабатывается отдельно.
_SVG_START_RE = re.compile(r"<svg\b", re.IGNORECASE)

# Lazy-parser: lxml с recover=True сам залечивает оборванные теги/кавычки.
_RECOVER_PARSER = etree.XMLParser(recover=True, huge_tree=True)


# ---------------------------------------------------------------------------
# Публичное API
# ---------------------------------------------------------------------------


def parse_to_valid(prefix: str, *, ensure_root: bool = True) -> str:
    """Превращает произвольный prefix в валидный SVG.

    Args:
        prefix: сырая строка (возможно, неполная, оборванная, с рассуждениями
            или markdown-обёрткой).
        ensure_root: если True (default), результат всегда начинается с
            `<svg ...>` и заканчивается `</svg>`. Если False — может вернуть
            фрагмент (для специальных случаев; в т.ч. пустую строку, если
            в prefix нет `<svg>`).

    Returns:
        Валидный SVG, который парсится через `lxml.etree.fromstring` без
        исключений. Внутри используется `lxml` для парсинга и `re` для
        препроцессинга.
    """
    # Защитное приведение типов: None → "", нестроковые → str().
    if prefix is None:
        prefix = ""
    if not isinstance(prefix, str):
        prefix = str(prefix)

    # 1. Найти первое `<svg` (case-insensitive). Если его нет — вернуть дефолт.
    m = _SVG_START_RE.search(prefix)
    if not m:
        return DEFAULT_EMPTY_SVG if ensure_root else ""

    # 2. Взять всё от первого `<svg`. Это отсекает reasoning, markdown-fence,
    #    "noise <svg ..." преамбулы и т.п.
    text = prefix[m.start():]

    # 3. Strip XML prolog / DOCTYPE, если они по какой-то причине попали сюда.
    text = _strip_prolog(text)

    # 4. Обрезать по первому `</svg>` (схлопывает trailing garbage, повторные
    #    SVG-блоки, невалидный хвост после валидного SVG).
    first_close = text.find("</svg")
    if first_close >= 0:
        text = text[: first_close + len("</svg>")]

    # 5. Попробовать распарсить в режиме recover. lxml сам залечит:
    #    - необструганный `<svg ...` (закроет как self-closing);
    #    - необструганный `<circle ...` (допишет `/>`);
    #    - открытую кавычку (`<rect fill="` → `<rect fill=""/>`);
    #    - атрибут `attr=` без значения (отбросит или сделает `attr=""`);
    #    - вложенные необструганные теги (закроет все уровни).
    root = _safe_parse(text)

    if root is None:
        # Совсем сломанный ввод (recover не справился). Возвращаем дефолт.
        return DEFAULT_EMPTY_SVG if ensure_root else ""

    if not ensure_root:
        return etree.tostring(root, encoding="unicode")

    # 6. Добавить xmlns, если отсутствует.
    root = _ensure_xmlns(root)

    # 7. Убедиться, что root — `<svg>`. Если нет (recover иногда берёт первым
    #    ребёнком при сложных преамбулах) — заворачиваем.
    if _local(root.tag) != "svg":
        root = _wrap_in_svg(root)
        root = _ensure_xmlns(root)

    return etree.tostring(root, encoding="unicode")


def stream_to_valid(chunks: Iterable[str]) -> Iterator[str]:
    """Принимает Iterable стримовых чанков, возвращает Iterator валидных
    SVG-snapshots после каждого нового чанка.

    Каждый snapshot — валидный SVG, пригодный для рендера.

    Example:
        >>> list(stream_to_valid(["<svg", ' viewBox="0 0 64 64"', ">", "</svg>"]))
        ['<svg xmlns="..."></svg>',
         '<svg xmlns="..." viewBox="0 0 64 64"></svg>',
         '<svg xmlns="..." viewBox="0 0 64 64"></svg>',
         '<svg xmlns="..." viewBox="0 0 64 64"></svg>']
    """
    buffer = ""
    for chunk in chunks:
        buffer += chunk
        yield parse_to_valid(buffer)


# ---------------------------------------------------------------------------
# Внутренние хелперы
# ---------------------------------------------------------------------------


def _strip_prolog(text: str) -> str:
    """Удаляет XML-декларацию `<?xml ...?>` и `<!DOCTYPE ...>` (если они есть)."""
    text = re.sub(r"<\?xml[^>]*\?>", "", text)
    text = re.sub(r"<!DOCTYPE[^>]*>", "", text)
    return text


def _safe_parse(text: str) -> etree._Element | None:
    """Парсит text в режиме recover. Возвращает root или None при полном фейле.

    recover=True сам закрывает необструганные теги/кавычки/атрибуты. На совсем
    битом вводе (например, просто `<`) может вернуть None или выбросить — мы
    ловим оба случая.
    """
    try:
        return etree.fromstring(text.encode("utf-8"), parser=_RECOVER_PARSER)
    except etree.XMLSyntaxError:
        return None
    except Exception:  # noqa: BLE001
        return None


def _ensure_xmlns(root: etree._Element) -> etree._Element:
    """Гарантирует, что у root есть `xmlns="http://www.w3.org/2000/svg"`.

    Если xmlns есть — возвращает root как есть. Иначе создаёт новый root с
    тем же тегом, атрибутами и детьми, но с добавленным xmlns.
    """
    if root.get("xmlns"):
        return root
    new_root = etree.Element(_local(root.tag), nsmap={None: SVG_NS})
    for k, v in root.attrib.items():
        new_root.set(_local(k), v)
    for child in root:
        new_root.append(child)
    return new_root


def _wrap_in_svg(root: etree._Element) -> etree._Element:
    """Заворачивает произвольный root в `<svg xmlns="...">...</svg>`."""
    new_root = etree.Element("svg", nsmap={None: SVG_NS})
    for k, v in root.attrib.items():
        new_root.set(_local(k), v)
    for child in root:
        new_root.append(child)
    return new_root


def _local(tag: object) -> str:
    """Снимает namespace-префикс: `{ns}local` → `local`. Терпимо к non-str."""
    s = str(tag)
    if "}" in s:
        return s.split("}", 1)[1]
    return s


__all__ = [
    "parse_to_valid",
    "stream_to_valid",
    "SVG_NS",
    "DEFAULT_EMPTY_SVG",
]