down2.py

#!/usr/bin/env python
# coding: utf-8

from __future__ import annotations

import atexit
import io
import json
import logging
import os
import pathlib
import re
import tempfile
import urllib.error
import urllib.parse
import urllib.request

from PIL import Image, ImageChops
from playwright.sync_api import Browser, sync_playwright
from pylatexenc.latex2text import LatexNodes2Text

import k3proc
from .mime import mime_to_suffix, mimetypes
from .syntax_highlight import code_to_html


logger = logging.getLogger(__name__)

_playwright = None
_browser = None


def _get_browser() -> Browser:
    global _playwright, _browser
    if _browser is None:
        _playwright = sync_playwright().start()
        _browser = _playwright.chromium.launch()
        atexit.register(_shutdown_browser)
    return _browser


def _shutdown_browser() -> None:
    global _playwright, _browser
    if _browser is not None:
        _browser.close()
        _browser = None
    if _playwright is not None:
        _playwright.stop()
        _playwright = None


zhihu_equation_url_fmt = "https://www.zhihu.com/equation?tex={texurl}{align}"

zhihu_equation_fmt = (
    '<img src="https://www.zhihu.com/equation'
    '?tex={texurl}{align}"'
    ' alt="{tex}{altalign}"'
    ' class="ee_img'
    ' tr_noresize"'
    ' eeimg="1">'
)


def convert(input_typ: str, content: str | bytes, output_typ: str, opt: dict[str, dict] | None = None) -> str | bytes:
    conv = mappings.get((input_typ, output_typ))
    if conv is None:
        raise ValueError(f"unsupported conversion: {input_typ} -> {output_typ}")
    if callable(conv):
        kwargs = {}
        if opt is not None and input_typ in opt:
            kwargs = opt[input_typ]
        return conv(content, **kwargs)
    else:
        #  indirect convertion
        inp = convert(input_typ, content, conv, opt=opt)
        return convert(conv, inp, output_typ, opt=opt)


def tex_to_zhihu_compatible(tex: str) -> tuple[str, str]:
    r"""
    Convert tex to zhihu compatible format.
    - ``>`` in img alt mess up the next escaped brace: ``\{ q > 1 \} --> \{ q > 1 }``.
    """

    tex = re.sub(r"\n", "", tex)
    tex = re.sub(r"(?<!\\)>", r"\\gt", tex)
    texurl = urllib.parse.quote(tex)
    return tex, texurl


def _zhihu_tex_params(tex: str, block: bool) -> dict[str, str]:
    """Prepare shared parameters for zhihu equation formatting."""
    tex, texurl = tex_to_zhihu_compatible(tex)
    # zhihu uses double backslash to center-align a block equation.
    align = "%5C%5C" if block else ""
    altalign = "\\\\" if block else ""
    return {"tex": tex, "texurl": texurl, "align": align, "altalign": altalign}


def tex_to_zhihu_url(tex: str, block: bool) -> str:
    """
    Convert tex source to a url linking to a svg on zhihu.
    www.zhihu.com/equation is a public api to render tex into svg.

    Args:
        tex(str): tex source

        block(bool): whether to render a block(center-aligned) equation or
            inline equation.

    Returns:
        string of a ``<img>`` tag.
    """

    p = _zhihu_tex_params(tex, block)
    return zhihu_equation_url_fmt.format(texurl=p["texurl"], align=p["align"])


def tex_to_zhihu(tex: str, block: bool) -> str:
    """
    Convert tex source to a img tag link to a svg on zhihu.
    www.zhihu.com/equation is a public api to render tex into svg.

    Args:
        tex(str): tex source

        block(bool): whether to render a block(center-aligned) equation or
            inline equation.

    Returns:
        string of a ``<img>`` tag.
    """

    p = _zhihu_tex_params(tex, block)
    return zhihu_equation_fmt.format(**p)


superscripts = {
    "0": "⁰",
    "1": "¹",
    "2": "²",
    "3": "³",
    "4": "⁴",
    "5": "⁵",
    "6": "⁶",
    "7": "⁷",
    "8": "⁸",
    "9": "⁹",
    "a": "ᵃ",
    "b": "ᵇ",
    "c": "ᶜ",
    "d": "ᵈ",
    "e": "ᵉ",
    "f": "ᶠ",
    "g": "ᵍ",
    "h": "ʰ",
    "i": "ⁱ",
    "j": "ʲ",
    "k": "ᵏ",
    "l": "ˡ",
    "m": "ᵐ",
    "n": "ⁿ",
    "o": "ᵒ",
    "p": "ᵖ",
    "r": "ʳ",
    "s": "ˢ",
    "t": "ᵗ",
    "u": "ᵘ",
    "v": "ᵛ",
    "w": "ʷ",
    "x": "ˣ",
    "y": "ʸ",
    "z": "ᶻ",
    "A": "ᴬ",
    "B": "ᴮ",
    "D": "ᴰ",
    "E": "ᴱ",
    "G": "ᴳ",
    "H": "ᴴ",
    "I": "ᴵ",
    "J": "ᴶ",
    "K": "ᴷ",
    "L": "ᴸ",
    "M": "ᴹ",
    "N": "ᴺ",
    "O": "ᴼ",
    "P": "ᴾ",
    "R": "ᴿ",
    "T": "ᵀ",
    "U": "ᵁ",
    "V": "ⱽ",
    "W": "ᵂ",
    "+": "⁺",
    "-": "⁻",
    "=": "⁼",
    "(": "⁽",
    ")": "⁾",
}

subscripts = {
    "0": "₀",
    "1": "₁",
    "2": "₂",
    "3": "₃",
    "4": "₄",
    "5": "₅",
    "6": "₆",
    "7": "₇",
    "8": "₈",
    "9": "₉",
    "+": "₊",
    "-": "₋",
    "=": "₌",
    "(": "₍",
    ")": "₎",
    "b": "ᵦ",
    "g": "ᵧ",
    "a": "ₐ",
    "e": "ₑ",
    "i": "ᵢ",
    "j": "ⱼ",
    "o": "ₒ",
    "r": "ᵣ",
    "u": "ᵤ",
    "v": "ᵥ",
    "x": "ₓ",
}


def all_in(chars: str, cate: dict[str, str]) -> bool:
    for c in chars:
        if c not in cate:
            return False
    return True


def tex_to_plain(tex: str) -> str:
    """
    Try hard converting tex to unicode plain text.
    """

    for reg, cate in (
        (r"_\{([^}]*?)\}", subscripts),
        (r"[\^]\{([^}]*?)\}", superscripts),
        (r"_(.)", subscripts),
        (r"[\^](.)", superscripts),
    ):
        pieces = []
        while True:
            match = re.search(reg, tex, flags=re.DOTALL | re.UNICODE)
            if match:
                chars = match.groups()[0]
                if all_in(chars, cate):
                    chars = [cate[x] for x in chars]
                else:
                    chars = tex[match.start() : match.end()]
                pieces.append(tex[: match.start()])
                pieces.append("".join(chars))
                tex = tex[match.end() :]
            else:
                pieces.append(tex)
                break

        tex = "".join(pieces)

    return LatexNodes2Text().latex_to_text(tex)


def tex_to_img(tex: str, block: bool, typ: str) -> str | bytes:
    """
    Convert tex source to an image.

    Args:
        tex(str): tex source

        block(bool): whether to render a block(center-aligned) equation or
            inline equation.

        typ(str): output image type such as "png" or "jpg"

    Returns:
        bytes of png data.
    """

    input_type = "tex_block"
    if not block:
        input_type = "tex_inline"
    return convert(input_type, tex, typ)


def download(url: str) -> bytes:
    """
    Download content from ``url`` and return the responded data.

    Args:
        url(str): the url from which to download.

    Returns:
        bytes of downloaded data.
    """

    resp = urllib.request.urlopen(url, timeout=30)
    return resp.read()


def web_to_img(pagefn: str, typ: str) -> bytes:
    """
    Render a web page, which could be html, svg etc into image.

    Args:
        pagefn(string): path to a local file that can be rendered in a browser.

        typ(string): specify output image type such as "png", "jpg"

    Returns:
        bytes of the image data
    """

    intyp = pagefn.rsplit(".")[-1]
    page = fread(pagefn)
    return render_to_img(intyp, page, typ)


def _trim_and_convert(png_data: bytes, typ: str) -> bytes:
    """Trim whitespace borders from a screenshot and convert to the target format."""
    img = Image.open(io.BytesIO(png_data))

    # Trim borders matching the corner pixel color (like ImageMagick -trim).
    bg = Image.new(img.mode, img.size, img.getpixel((0, 0)))
    diff = ImageChops.difference(img, bg)
    bbox = diff.getbbox()
    if bbox:
        img = img.crop(bbox)

    buf = io.BytesIO()
    if typ == "png":
        img.save(buf, format="PNG")
    else:
        background = Image.new("RGB", img.size, (255, 255, 255))
        if img.mode == "RGBA":
            background.paste(img, mask=img.split()[3])
        else:
            background.paste(img)
        background.save(buf, format="JPEG")

    return buf.getvalue()


def render_to_img(
    mime: str, content: str | bytes, typ: str, width: int = 1000, height: int = 2000, asset_base: str | None = None
) -> bytes:
    """
    Render content that is renderable in a browser to image.
    Such as html, svg etc into image.
    Uses Playwright (Chromium) for rendering and Pillow for image processing.

    Args:
        mime(str): a full mime type such as ``image/jpeg`` or a shortcut ``jpg``.

        content(str): content to render, such as jpeg data or svg source.

        typ(string): specifies output image type such as "png", "jpg"

        width(int): specifies the window width to render a page. Default 1000.

        height(int): specifies the window height to render a page. Default 2000.

        asset_base(str): specifies the path to assets dir. E.g. the image base path in a html page.

    Returns:
        bytes of the image data
    """

    if "html" in mime:
        content = r'<meta http-equiv="Content-Type" content="text/html; charset=utf-8"/>' + content

        if asset_base is not None:
            base_uri = pathlib.Path(asset_base).as_uri()
            content = '<base href="{}/">'.format(base_uri) + content

    m = mimetypes.get(mime) or mime
    suffix = mime_to_suffix.get(m, mime)

    with tempfile.TemporaryDirectory() as tdir:
        fn = os.path.join(tdir, "xxx." + suffix)
        flags = "w"
        if isinstance(content, bytes):
            flags = "wb"
        with open(fn, flags) as f:
            f.write(content)

        browser = _get_browser()
        page = browser.new_page(
            viewport={"width": width, "height": height},
            device_scale_factor=2,
        )
        page.goto(pathlib.Path(fn).as_uri())

        content_height = page.evaluate("document.documentElement.scrollHeight")
        if content_height > height:
            page.set_viewport_size({"width": width, "height": content_height})

        png_data = page.screenshot(omit_background=True)
        page.close()

    return _trim_and_convert(png_data, typ)


html_style = """
<style type="text/css" media="screen">
    table {
        display: block;
        margin-bottom: 1em;
        width: fit-content;
        font-family: -apple-system,BlinkMacSystemFont,"Roboto","Segoe UI","Helvetica Neue","Lucida Grande",Arial,sans-serif;
        font-size: .75em;
        border-collapse: collapse;
        overflow-x: auto;
    }

    thead {
        background-color: #f2f3f3;
        border-bottom: 2px solid #b6b6b6;
    }

    th {
        padding: 0.5em;
        font-weight: bold;
        text-align: left;
    }

    td {
        padding: 0.5em;
        border-bottom: 1px solid #ddd;
    }

    tr,
    td,
    th {
        vertical-align: middle;
    }
    pre.highlight {
        margin: 0;
        padding: 1em;
        background: #263238;
        color: #eff;
        font-size: 1.5em;
        font-family: "SFMono-Regular",Consolas,"Liberation Mono",Menlo,Courier,"PingFang SC", "Microsoft YaHei",monospace;
    }
</style>
"""


def md_to_html(md: str) -> str:
    """
    Build markdown source into html.

    Args:
        md(str): markdown source.

    Returns:
        str of html
    """

    _, html, _ = k3proc.command_ex(
        "pandoc",
        "-f",
        "markdown",
        "-t",
        "html",
        input=md,
    )

    return html_style + html


def mdtable_to_barehtml(md: str) -> str:
    """
    Build markdown table into html without style.

    Args:
        md(str): markdown source.

    Returns:
        str of html
    """

    # A table with wide column will cause pandoc to produce ``colgroup`` tag, which is not recognized by zhihu.
    # Reported in:
    #      https://github.com/drmingdrmer/md2zhihu/issues/22
    #
    # Thus we have to set a very big rendering window to disable this behavior
    #      https://github.com/jgm/pandoc/issues/2574

    _, html, _ = k3proc.command_ex(
        "pandoc",
        "-f",
        "markdown",
        "-t",
        "html",
        "--column",
        "100000",
        input=md,
    )
    lines = html.strip().split("\n")
    lines = [x for x in lines if x not in ("<thead>", "</thead>", "<tbody>", "</tbody>")]

    return "\n".join(lines)


def mermaid_to_svg(mmd: str) -> str:
    """
    Render mermaid to svg.
    See: https://mermaid-js.github.io/mermaid/#

    Requires:
        npm install @mermaid-js/mermaid-cli
    """

    with tempfile.TemporaryDirectory() as tdir:
        output_path = os.path.join(tdir, "mmd.svg")

        puppeteer_config = {
            "args": [
                "--no-sandbox",
                "--disable-setuid-sandbox",
                "--disable-dev-shm-usage",
                "--disable-gpu",
            ]
        }

        config_file_path = os.path.join(tdir, "config.json")
        with open(config_file_path, "w") as f:
            f.write(json.dumps(puppeteer_config))

        k3proc.command_ex(
            "npm",
            "exec",
            "--",
            "mmdc",
            "-o",
            output_path,
            "--puppeteerConfigFile",
            config_file_path,
            input=mmd,
        )
        return fread(output_path)


def graphviz_to_img(gv: str | bytes, typ: str) -> bytes:
    """
    Render graphviz source to image.

    Requires:
        brew install graphviz
    """

    _, out, _ = k3proc.command_ex(
        "dot",
        "-T" + typ,
        input=to_bytes(gv),
        text=False,
    )
    return out


def to_bytes(s: str | bytes) -> bytes:
    if isinstance(s, bytes):
        return s
    return bytes(s, "utf-8")


def pjoin(*p: str) -> str:
    return os.path.join(*p)


def fread(*p: str) -> str:
    with open(os.path.join(*p), "r") as f:
        return f.read()


mappings = {
    ("md", "html"): md_to_html,
    ("md", "jpg"): "html",
    ("md", "png"): "html",
    ("html", "jpg"): lambda x, **kwargs: render_to_img("html", x, "jpg", **kwargs),
    ("html", "png"): lambda x, **kwargs: render_to_img("html", x, "png", **kwargs),
    # markdown table
    ("table", "html"): mdtable_to_barehtml,
    ("table", "jpg"): "html",
    ("table", "png"): "html",
    ("mermaid", "svg"): mermaid_to_svg,
    ("mermaid", "jpg"): "svg",
    ("mermaid", "png"): "svg",
    ("graphviz", "svg"): lambda x: graphviz_to_img(x, "svg"),
    ("graphviz", "jpg"): lambda x: graphviz_to_img(x, "jpg"),
    ("graphviz", "png"): lambda x: graphviz_to_img(x, "png"),
    ("tex_block", "url"): lambda x: tex_to_zhihu_url(x, True),
    ("tex_inline", "url"): lambda x: tex_to_zhihu_url(x, False),
    ("tex_block", "imgtag"): lambda x: tex_to_zhihu(x, True),
    ("tex_inline", "imgtag"): lambda x: tex_to_zhihu(x, False),
    ("url", "jpg"): download,
    ("url", "png"): download,
    ("url", "svg"): download,
    ("url", "html"): download,
    ("tex_block", "jpg"): "svg",
    ("tex_inline", "jpg"): "svg",
    ("tex_block", "png"): "svg",
    ("tex_inline", "png"): "svg",
    ("tex_block", "svg"): "url",
    ("tex_inline", "svg"): "url",
    ("tex_inline", "plain"): tex_to_plain,
    ("svg", "jpg"): lambda x: render_to_img("svg", x, "jpg"),
    ("svg", "png"): lambda x: render_to_img("svg", x, "png"),
    ("code", "html"): code_to_html,
    ("code", "jpg"): "html",
    ("code", "png"): "html",
}