Store document images locally

#!/usr/bin/env python3
import argparse
import html
import re
import time
import urllib.error
import urllib.parse
import urllib.request
from io import BytesIO
from pathlib import Path

from PIL import Image, ImageOps, UnidentifiedImageError

MARKDOWN_IMAGE_RE = re.compile(
    r"!\[[^\]\r\n]*\]\(\s*(?:<(?P<angle_url>https?://[^>\r\n]+)>|"
    r"(?P<plain_url>https?://[^\s)\r\n]+))",
    re.IGNORECASE,
)
HTML_IMAGE_RE = re.compile(
    r"<img\b[^>]*?\bsrc\s*=\s*(?P<quote>[\"'])(?P<url>https?://.*?)(?P=quote)",
    re.IGNORECASE | re.DOTALL,
)
RETRYABLE_STATUS = {408, 425, 429, 500, 502, 503, 504}


def find_urls(markdown):
    spans = []
    for match in MARKDOWN_IMAGE_RE.finditer(markdown):
        group = "angle_url" if match.group("angle_url") else "plain_url"
        spans.append(
            (match.start(group), match.end(group), html.unescape(match.group(group)))
        )
    for match in HTML_IMAGE_RE.finditer(markdown):
        spans.append(
            (match.start("url"), match.end("url"), html.unescape(match.group("url")))
        )
    spans.sort(key=lambda item: item[0])

    result = []
    previous_end = -1
    for span in spans:
        if span[0] >= previous_end:
            result.append(span)
            previous_end = span[1]
    return result


def download(url, timeout, retries):
    request = urllib.request.Request(
        url,
        headers={
            "Accept": "image/*,*/*;q=0.8",
            "User-Agent": "somark-image-example/1.0",
        },
    )
    for attempt in range(retries + 1):
        try:
            with urllib.request.urlopen(request, timeout=timeout) as response:
                content_type = response.headers.get_content_type().lower()
                if not (
                    content_type.startswith("image/")
                    or content_type == "application/octet-stream"
                ):
                    raise ValueError(
                        f"URL returned non-image content: {content_type}"
                    )
                return response.read()
        except urllib.error.HTTPError as error:
            if error.code not in RETRYABLE_STATUS or attempt == retries:
                raise
        except (urllib.error.URLError, TimeoutError, OSError):
            if attempt == retries:
                raise
        time.sleep(min(2**attempt, 8))
    raise RuntimeError("unreachable")


def to_jpeg(data, destination):
    try:
        with Image.open(BytesIO(data)) as opened:
            opened.seek(0)
            oriented = ImageOps.exif_transpose(opened)
            try:
                oriented.load()
                has_alpha = oriented.mode in {"RGBA", "LA"} or (
                    oriented.mode == "P" and "transparency" in oriented.info
                )
                if has_alpha:
                    rgba = oriented.convert("RGBA")
                    converted = Image.new("RGB", rgba.size, (255, 255, 255))
                    converted.paste(rgba, mask=rgba.getchannel("A"))
                    rgba.close()
                else:
                    converted = oriented.convert("RGB")
                try:
                    converted.save(
                        destination,
                        "JPEG",
                        quality=95,
                        subsampling=0,
                        optimize=True,
                    )
                finally:
                    converted.close()
            finally:
                if oriented is not opened:
                    oriented.close()
    except (UnidentifiedImageError, OSError, ValueError) as error:
        raise RuntimeError(f"Invalid image: {error}") from error


def rewrite(markdown, spans, names):
    chunks = []
    cursor = 0
    for start, end, url in spans:
        chunks.append(markdown[cursor:start])
        chunks.append(
            urllib.parse.quote("./images/" + names[url], safe="/-._~")
        )
        cursor = end
    chunks.append(markdown[cursor:])
    return "".join(chunks)


def localize(input_path, output_dir, timeout, retries):
    input_path = input_path.resolve()
    output_dir = output_dir.resolve()
    if (
        input_path.suffix.lower() not in {".md", ".markdown"}
        or not input_path.is_file()
    ):
        raise ValueError("Input must be an existing .md or .markdown file")
    if output_dir.exists():
        raise FileExistsError(f"Output directory already exists: {output_dir}")

    markdown = input_path.read_text(encoding="utf-8-sig")
    spans = find_urls(markdown)
    urls = list(dict.fromkeys(span[2] for span in spans))
    names = {
        url: f"image_{index:03d}.jpg"
        for index, url in enumerate(urls, start=1)
    }

    downloaded = {url: download(url, timeout, retries) for url in urls}
    output_dir.mkdir(parents=True)
    image_dir = output_dir / "images"
    image_dir.mkdir()
    try:
        for url in urls:
            to_jpeg(downloaded[url], image_dir / names[url])
        with (output_dir / "main.md").open(
            "w", encoding="utf-8", newline=""
        ) as output:
            output.write(rewrite(markdown, spans, names))
    except Exception:
        for path in image_dir.glob("image_*.jpg"):
            path.unlink()
        raise


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input", type=Path)
    parser.add_argument("-o", "--output-dir", type=Path)
    parser.add_argument("--timeout", type=float, default=60.0)
    parser.add_argument("--retries", type=int, default=3)
    args = parser.parse_args()
    output_dir = args.output_dir or args.input.with_suffix("")
    localize(args.input, output_dir, args.timeout, args.retries)
    print(f"Completed: {output_dir / 'main.md'}")


if __name__ == "__main__":
    main()