browser_qisi_scraper.py

# -*- coding: utf-8 -*-
"""
Browser-backed Xueqiu scraper.

This uses Playwright with a persistent browser profile. You log in once in the
opened browser; the script then calls Xueqiu's timeline API from the page
context, so the requests carry the browser's own cookies.
"""

from __future__ import annotations

import argparse
import json
import random
import sqlite3
import time
from datetime import datetime
from pathlib import Path
from typing import Any

from direct_qisi_scraper import (
    BASE_URL,
    DEFAULT_DOMAIN,
    DEFAULT_UID,
    init_output_db,
    normalize_status,
    parse_snowman_status,
    save_sqlite,
    write_jsonl,
    write_markdown,
)


def import_playwright() -> Any:
    try:
        from playwright.sync_api import TimeoutError as PlaywrightTimeoutError
        from playwright.sync_api import sync_playwright
    except ImportError as exc:
        raise SystemExit(
            "缺少 Playwright。先运行：\n"
            "  python -m pip install playwright\n"
            "通常使用本机 Chrome/Edge 不需要下载 Playwright Chromium。"
        ) from exc
    return sync_playwright, PlaywrightTimeoutError


def evaluate_fetch(page: Any, path: str) -> dict[str, Any]:
    return page.evaluate(
        """
        async (path) => {
          const resp = await fetch(path, {
            credentials: 'include',
            headers: {
              'accept': 'application/json, text/plain, */*',
              'x-requested-with': 'XMLHttpRequest'
            }
          });
          const text = await resp.text();
          let data = null;
          try { data = JSON.parse(text); } catch (e) {}
          return {ok: resp.ok, status: resp.status, statusText: resp.statusText, text, data};
        }
        """,
        path,
    )


def timeline_path(uid: str, page_num: int, status_type: str = "") -> str:
    path = f"/v4/statuses/user_timeline.json?user_id={uid}&page={page_num}"
    if status_type:
        path += f"&type={status_type}"
    return path


def wait_for_login(page: Any, uid: str, timeout_seconds: int) -> bool:
    deadline = time.time() + timeout_seconds
    while time.time() < deadline:
        result = evaluate_fetch(page, timeline_path(uid, 2))
        if result["ok"]:
            return True
        time.sleep(3)
    return False


def fetch_article_from_browser(page: Any, target: str) -> dict[str, Any] | None:
    result = page.evaluate(
        """
        async (target) => {
          const resp = await fetch(target, {credentials: 'include', headers: {'accept': 'text/html'}});
          return {ok: resp.ok, status: resp.status, text: await resp.text()};
        }
        """,
        target,
    )
    if not result["ok"]:
        return None
    return parse_snowman_status(result["text"])


def existing_rows(db_path: Path) -> dict[int, dict[str, Any]]:
    if not db_path.exists():
        return {}
    conn = sqlite3.connect(db_path)
    try:
        rows = conn.execute("SELECT raw_json, full_raw_json FROM statuses").fetchall()
    except sqlite3.Error:
        conn.close()
        return {}
    conn.close()
    parsed = {}
    for raw_json, full_raw_json in rows:
        raw = json.loads(raw_json)
        full = json.loads(full_raw_json) if full_raw_json else None
        sid = raw.get("id") or (full or {}).get("id")
        if sid:
            parsed[int(sid)] = normalize_status(raw, full)
    return parsed


def build_output_dir(base: str, resume_dir: str) -> Path:
    if resume_dir:
        out_dir = Path(resume_dir)
        out_dir.mkdir(parents=True, exist_ok=True)
        return out_dir
    out_dir = Path(base) / f"qisi_browser_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    out_dir.mkdir(parents=True, exist_ok=True)
    return out_dir


def choose_channel(value: str) -> str | None:
    return value if value and value != "bundled" else None


def crawl(args: argparse.Namespace) -> int:
    sync_playwright, PlaywrightTimeoutError = import_playwright()
    out_dir = build_output_dir(args.output_dir, args.resume_dir)
    profile_dir = Path(args.profile_dir).expanduser()
    rows_by_id = existing_rows(out_dir / "timeline.sqlite") if args.resume_dir else {}
    raw_pages = []
    requested_all = args.pages == "all" and args.max_pages is None

    with sync_playwright() as p:
        browser_type = p.chromium
        launch_kwargs: dict[str, Any] = {
            "headless": args.headless,
            "args": ["--disable-blink-features=AutomationControlled"],
        }
        channel = choose_channel(args.channel)
        if channel:
            launch_kwargs["channel"] = channel
        context = browser_type.launch_persistent_context(str(profile_dir), **launch_kwargs)
        page = context.pages[0] if context.pages else context.new_page()
        page.goto(f"{BASE_URL}/{args.domain}", wait_until="domcontentloaded", timeout=60000)

        print(f"[browser] profile={profile_dir}")
        print(f"[browser] target={BASE_URL}/{args.domain}")
        if not args.headless:
            print("[browser] 如浏览器未登录雪球，请在打开的窗口完成登录。")

        login_ok = wait_for_login(page, args.uid, args.login_timeout)
        if not login_ok:
            context.close()
            print("[browser] 登录态验证失败：第 2 页仍不可访问。请确认浏览器已登录雪球。")
            return 2

        first = evaluate_fetch(page, timeline_path(args.uid, 1, args.type))
        if not first["ok"] or not first["data"]:
            context.close()
            print(f"[browser] 第 1 页获取失败：HTTP {first['status']} {first['text'][:200]}")
            return 2

        advertised_max_page = int(first["data"].get("maxPage") or 1)
        advertised_total = int(first["data"].get("total") or first["data"].get("count") or 0)
        max_page = advertised_max_page
        if args.pages and args.pages != "all":
            max_page = min(max_page, int(args.pages))
        if args.max_pages:
            max_page = min(max_page, args.max_pages)

        stop_reason = "page_limit" if max_page < advertised_max_page else "completed"
        failed_page = None
        print(f"[browser] max_page={max_page}")
        for page_num in range(1, max_page + 1):
            result = first if page_num == 1 else evaluate_fetch(page, timeline_path(args.uid, page_num, args.type))
            if not result["ok"] or not result["data"]:
                failed_page = page_num
                stop_reason = f"http_{result['status']}"
                print(f"[browser] page {page_num} failed: HTTP {result['status']} {result['text'][:200]}")
                if args.stop_on_error:
                    context.close()
                    return 2
                break

            data = result["data"]
            statuses = data.get("statuses") or data.get("list") or []
            raw_pages.append({"page": page_num, "response": data})
            print(f"[browser] page {page_num}/{max_page}: {len(statuses)} statuses")

            for raw in statuses:
                sid = raw.get("id")
                if not sid:
                    continue
                full = None
                if args.fetch_articles and (raw.get("is_column") or str(raw.get("type")) == "3"):
                    target = raw.get("target") or f"/{args.uid}/{sid}"
                    try:
                        full = fetch_article_from_browser(page, target)
                        time.sleep(max(0, args.sleep + random.uniform(0, args.jitter)))
                    except (PlaywrightTimeoutError, Exception) as exc:
                        print(f"[browser] article {sid} fetch failed: {exc}")
                rows_by_id[int(sid)] = normalize_status(raw, full)

            if page_num % args.flush_every == 0:
                flush_outputs(
                    out_dir,
                    args.uid,
                    rows_by_id,
                    raw_pages,
                    cookie_source="browser",
                    advertised_total=advertised_total,
                    advertised_max_page=advertised_max_page,
                    requested_max_page=max_page,
                    requested_all=requested_all,
                    stop_reason="checkpoint",
                    failed_page=None,
                )
            time.sleep(max(0, args.sleep + random.uniform(0, args.jitter)))

        context.close()

    flush_outputs(
        out_dir,
        args.uid,
        rows_by_id,
        raw_pages,
        cookie_source="browser",
        advertised_total=advertised_total,
        advertised_max_page=advertised_max_page,
        requested_max_page=max_page,
        requested_all=requested_all,
        stop_reason=stop_reason,
        failed_page=failed_page,
    )
    print(f"[browser] saved {len(rows_by_id)} statuses to {out_dir}")
    return 0


def flush_outputs(
    out_dir: Path,
    uid: str,
    rows_by_id: dict[int, dict[str, Any]],
    raw_pages: list[dict[str, Any]],
    cookie_source: str,
    advertised_total: int | None,
    advertised_max_page: int | None,
    requested_max_page: int | None,
    requested_all: bool,
    stop_reason: str,
    failed_page: int | None,
) -> None:
    rows = sorted(rows_by_id.values(), key=lambda item: item.get("created_at") or "", reverse=True)
    (out_dir / "timeline_raw_pages.json").write_text(json.dumps(raw_pages, ensure_ascii=False, indent=2), encoding="utf-8")
    write_jsonl(out_dir / "timeline_clean.jsonl", rows)
    write_markdown(out_dir / "timeline_clean.md", rows, uid)
    conn = init_output_db(out_dir / "timeline.sqlite")
    save_sqlite(conn, rows)
    conn.close()
    summary = {
        "uid": uid,
        "status_count": len(rows),
        "column_count": sum(1 for row in rows if row.get("is_column")),
        "column_fulltext_count": sum(1 for row in rows if row.get("is_column") and row.get("full_raw")),
        "pages_saved": len(raw_pages),
        "advertised_total": advertised_total,
        "advertised_max_page": advertised_max_page,
        "requested_max_page": requested_max_page,
        "requested_all_pages": requested_all,
        "complete": bool(
            requested_all
            and advertised_max_page is not None
            and failed_page is None
            and len(raw_pages) >= advertised_max_page
        ),
        "stop_reason": stop_reason,
        "failed_page": failed_page,
        "output_dir": str(out_dir.resolve()),
        "generated_at": datetime.now().isoformat(timespec="seconds"),
        "cookie_source": cookie_source,
    }
    (out_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Scrape Xueqiu qisi timeline through a logged-in browser.")
    parser.add_argument("--uid", default=DEFAULT_UID)
    parser.add_argument("--domain", default=DEFAULT_DOMAIN)
    parser.add_argument("--output-dir", default="output")
    parser.add_argument("--resume-dir", default="", help="Existing output directory to update.")
    parser.add_argument("--profile-dir", default=str(Path.home() / ".xueqiu_playwright_profile"))
    parser.add_argument("--channel", default="chrome", help="chrome, msedge, or bundled")
    parser.add_argument("--headless", action="store_true")
    parser.add_argument("--pages", default="all", help="'all' or a page count")
    parser.add_argument("--max-pages", type=int, default=None)
    parser.add_argument("--type", default="")
    parser.add_argument("--fetch-articles", action="store_true")
    parser.add_argument("--login-timeout", type=int, default=300)
    parser.add_argument("--sleep", type=float, default=1.0)
    parser.add_argument("--jitter", type=float, default=0.7)
    parser.add_argument("--flush-every", type=int, default=10)
    parser.add_argument("--stop-on-error", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    raise SystemExit(crawl(parse_args()))