86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""HTML → Markdown via markdownify. Images only resolve local file refs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import re
|
|
from pathlib import Path
|
|
from urllib.parse import unquote, urlparse
|
|
|
|
from . import ParsedImage, ParsedPage, ParseResult
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_SKIP_TAGS = ("script", "style", "noscript", "template")
|
|
|
|
|
|
def _collect_local_images(soup, html_path: Path) -> list[ParsedImage]:
|
|
images: list[ParsedImage] = []
|
|
base_dir = html_path.parent
|
|
for tag in soup.find_all("img"):
|
|
src = tag.get("src") or ""
|
|
parsed = urlparse(src)
|
|
if parsed.scheme and parsed.scheme not in ("file", ""):
|
|
continue
|
|
local_path = (base_dir / unquote(parsed.path or src)).resolve()
|
|
try:
|
|
local_path.relative_to(base_dir.resolve())
|
|
except ValueError:
|
|
# Path traversal: refuse to read outside the source's directory.
|
|
continue
|
|
if not local_path.is_file():
|
|
continue
|
|
try:
|
|
blob = local_path.read_bytes()
|
|
except OSError:
|
|
continue
|
|
suffix = local_path.suffix.lower().lstrip(".")
|
|
mime = {
|
|
"png": "image/png",
|
|
"jpg": "image/jpeg",
|
|
"jpeg": "image/jpeg",
|
|
"gif": "image/gif",
|
|
"webp": "image/webp",
|
|
"svg": "image/svg+xml",
|
|
}.get(suffix, f"image/{suffix or 'octet-stream'}")
|
|
caption = tag.get("alt") or tag.get("title") or ""
|
|
images.append(
|
|
ParsedImage(
|
|
image_bytes = blob,
|
|
mime_type = mime,
|
|
page_number = None,
|
|
nearest_caption = caption,
|
|
)
|
|
)
|
|
return images
|
|
|
|
|
|
def extract(path: Path, *, want_images: bool = False) -> ParseResult:
|
|
from bs4 import BeautifulSoup
|
|
from markdownify import markdownify
|
|
|
|
raw = path.read_bytes()
|
|
soup = BeautifulSoup(raw, "lxml")
|
|
for tag_name in _SKIP_TAGS:
|
|
for tag in soup.find_all(tag_name):
|
|
tag.decompose()
|
|
|
|
images: list[ParsedImage] = []
|
|
if want_images:
|
|
images = _collect_local_images(soup, path)
|
|
|
|
md = markdownify(
|
|
str(soup),
|
|
heading_style = "ATX",
|
|
strip = list(_SKIP_TAGS),
|
|
)
|
|
md = re.sub(r"\n{3,}", "\n\n", md).strip()
|
|
if not md:
|
|
return ParseResult(pages = [], images = images)
|
|
return ParseResult(
|
|
pages = [ParsedPage(text = md, page_number = None)],
|
|
images = images,
|
|
)
|