html_to_image.py

#!/usr/bin/env python3
"""
将HTML文章转换为长图
"""

import os
import sys
import argparse
from pathlib import Path


def html_to_image_playwright(html_file: str, output_file: str, width: int = 800):
    """
    使用 Playwright 将 HTML 转换为长图

    Args:
        html_file: HTML 文件路径
        output_file: 输出图片路径
        width: 图片宽度（像素）
    """
    try:
        from playwright.sync_api import sync_playwright
    except ImportError:
        print("❌ 需要安装 playwright")
        print("   安装命令: pip install playwright")
        print("   然后运行: playwright install chromium")
        sys.exit(1)

    print(f"📄 读取 HTML: {html_file}")

    # 读取 HTML 内容
    with open(html_file, 'r', encoding='utf-8') as f:
        html_content = f.read()

    # 将图片路径转换为 file:// 协议
    html_dir = os.path.dirname(os.path.abspath(html_file))
    skill_dir = os.path.dirname(html_dir)  # 上一级目录（wechat-burst-gen）

    # 处理绝对路径的图片（已经包含完整路径）
    import re
    # 匹配 src="/Users/..." 格式的绝对路径
    html_content = re.sub(
        r'src="(/Users/[^"]+)"',
        r'src="file://\1"',
        html_content
    )
    html_content = re.sub(
        r"src='(/Users/[^']+)'",
        r"src='file://\1'",
        html_content
    )

    # 处理相对路径的图片
    html_content = html_content.replace('src="images/', f'src="file://{skill_dir}/images/')
    html_content = html_content.replace('src="../images/', f'src="file://{skill_dir}/images/')
    html_content = html_content.replace("src='images/", f"src='file://{skill_dir}/images/")
    html_content = html_content.replace("src='../images/", f"src='file://{skill_dir}/images/")

    with sync_playwright() as p:
        print("🌐 启动浏览器...")
        browser = p.chromium.launch()
        page = browser.new_page(viewport={'width': width, 'height': 800})

        # 加载 HTML
        print("📝 渲染页面...")
        page.set_content(html_content, wait_until='networkidle')

        # 等待所有图片加载（增加等待时间）
        print("⏳ 等待图片加载...")
        page.wait_for_timeout(5000)  # 增加到5秒

        # 等待所有图片标签加载完成
        try:
            page.wait_for_load_state('networkidle', timeout=10000)
        except:
            pass  # 忽略超时错误，继续截图

        # 获取页面完整高度
        height = page.evaluate("document.documentElement.scrollHeight")
        print(f"📐 页面尺寸: {width}x{height}px")

        # 截取完整页面
        print("📸 生成长图...")
        page.screenshot(path=output_file, full_page=True)

        browser.close()

    print(f"✅ 长图已保存: {output_file}")

    # 获取文件大小
    size_kb = os.path.getsize(output_file) / 1024
    print(f"📊 文件大小: {size_kb:.1f} KB")


def html_to_image_selenium(html_file: str, output_file: str, width: int = 800):
    """
    使用 Selenium 将 HTML 转换为长图（备选方案）

    Args:
        html_file: HTML 文件路径
        output_file: 输出图片路径
        width: 图片宽度（像素）
    """
    try:
        from selenium import webdriver
        from selenium.webdriver.chrome.options import Options
    except ImportError:
        print("❌ 需要安装 selenium")
        print("   安装命令: pip install selenium")
        sys.exit(1)

    print(f"📄 读取 HTML: {html_file}")

    # Chrome 选项
    chrome_options = Options()
    chrome_options.add_argument('--headless')
    chrome_options.add_argument('--disable-gpu')
    chrome_options.add_argument(f'--window-size={width},1080')

    # 启动浏览器
    print("🌐 启动浏览器...")
    driver = webdriver.Chrome(options=chrome_options)

    try:
        # 加载 HTML
        html_path = f"file://{os.path.abspath(html_file)}"
        driver.get(html_path)

        # 等待加载
        import time
        time.sleep(2)

        # 获取完整页面高度
        height = driver.execute_script("return document.documentElement.scrollHeight")

        # 设置窗口大小为完整页面大小
        driver.set_window_size(width, height)

        print(f"📐 页面尺寸: {width}x{height}px")
        print("📸 生成长图...")

        # 截图
        driver.save_screenshot(output_file)

        print(f"✅ 长图已保存: {output_file}")

        # 获取文件大小
        size_kb = os.path.getsize(output_file) / 1024
        print(f"📊 文件大小: {size_kb:.1f} KB")

    finally:
        driver.quit()


def main():
    parser = argparse.ArgumentParser(description='将 HTML 文章转换为长图')
    parser.add_argument('html_file', nargs='?', help='HTML 文件路径')
    parser.add_argument('-o', '--output', help='输出图片路径（默认：同名 .png）')
    parser.add_argument('-w', '--width', type=int, default=800, help='图片宽度（默认：800px）')
    parser.add_argument('--method', choices=['playwright', 'selenium'], default='playwright',
                       help='转换方法（默认：playwright）')

    args = parser.parse_args()

    # 确定要转换的文件
    if args.html_file:
        html_file = args.html_file
    else:
        # 查找最新的 HTML 文件
        output_dir = "./output"
        html_files = list(Path(output_dir).glob("*.html"))

        if not html_files:
            print("❌ 未找到 HTML 文件")
            print("   请指定 HTML 文件路径")
            sys.exit(1)

        # 按修改时间排序，取最新的
        html_file = str(sorted(html_files, key=lambda x: x.stat().st_mtime, reverse=True)[0])
        print(f"📄 使用最新文章: {html_file}\n")

    # 检查文件是否存在
    if not os.path.exists(html_file):
        print(f"❌ 文件不存在: {html_file}")
        sys.exit(1)

    # 确定输出文件名
    if args.output:
        output_file = args.output
    else:
        # 使用同名 .png
        base_name = os.path.splitext(html_file)[0]
        output_file = f"{base_name}_long_image.png"

    print("\n" + "=" * 70)
    print("📸 HTML 转长图工具")
    print("=" * 70 + "\n")

    # 转换
    try:
        if args.method == 'playwright':
            html_to_image_playwright(html_file, output_file, args.width)
        else:
            html_to_image_selenium(html_file, output_file, args.width)

        print("\n" + "=" * 70)
        print("🎉 转换完成！")
        print("=" * 70)
        print(f"\n💡 打开图片: open \"{output_file}\"")
        print()

    except Exception as e:
        print(f"\n❌ 转换失败: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()