예제 소스

sgf_download.py

"""u-go.net 의 KGS4d (4단 이상 KGS 기보) 월별 tar.gz 자동 다운로드.

URL 패턴:
    https://dl.u-go.net/gamerecords-4d/KGS4d-YYYY_MM-19-COUNT-.tar.gz

`COUNT` 는 매 월마다 다른 값이라 직접 지정할 수 없으므로,
인덱스 페이지(`https://u-go.net/gamerecords-4d/`) 를 한 번 긁어 (year, month) → URL 매핑을 만든 뒤
사용자가 지정한 기간만 다운로드합니다.

사용:
    python -m vlm_server.sgf_download --years 2018 --out-dir sgf_data
    python -m vlm_server.sgf_download --years 2017 2018 --out-dir sgf_data
    python -m vlm_server.sgf_download --months 2018-01 2018-02 --out-dir sgf_data
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import tarfile
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Iterable, List, Optional, Tuple

INDEX_URL = "https://u-go.net/gamerecords-4d/"
DL_URL_RE = re.compile(
    r"https?://dl\.u-go\.net/gamerecords-4d/KGS4d-(\d{4})_(\d{2})-19-(\d+)-\.tar\.gz",
    re.IGNORECASE,
)
DL_YEAR_URL_RE = re.compile(
    r"https?://dl\.u-go\.net/gamerecords-4d/KGS4d-(\d{4})-19-(\d+)-\.tar\.gz",
    re.IGNORECASE,
)


def fetch_index() -> List[Tuple[str, str, str]]:
    """(yyyymm, count, url) 리스트. 'yyyymm' = 연간 통합본은 'YYYY' 만."""
    req = urllib.request.Request(
        INDEX_URL,
        headers={"User-Agent": "baduk-vlm-server/0.1"},
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        body = resp.read().decode("utf-8", errors="replace")
    out: List[Tuple[str, str, str]] = []
    for m in DL_URL_RE.finditer(body):
        y, mo, cnt = m.group(1), m.group(2), m.group(3)
        out.append((f"{y}{mo}", cnt, m.group(0)))
    for m in DL_YEAR_URL_RE.finditer(body):
        y, cnt = m.group(1), m.group(2)
        out.append((y, cnt, m.group(0)))
    # 중복 제거 (월별/연별이 모두 있으면 둘 다 둠. 사용자가 골라 받음)
    seen = set()
    uniq: List[Tuple[str, str, str]] = []
    for k, c, u in out:
        key = (k, u)
        if key in seen:
            continue
        seen.add(key)
        uniq.append((k, c, u))
    return uniq


def parse_month_token(s: str) -> str:
    """`'2018-01'` → `'201801'`, `'2018'` → `'2018'`."""
    s = s.strip()
    if "-" in s:
        y, mo = s.split("-", 1)
        return f"{int(y):04d}{int(mo):02d}"
    return f"{int(s):04d}"


def select_archives(
    index: List[Tuple[str, str, str]],
    years: Optional[Iterable[str]] = None,
    months: Optional[Iterable[str]] = None,
    prefer_monthly: bool = True,
) -> List[Tuple[str, str, str]]:
    """필터링."""
    target_keys: set[str] = set()
    if years:
        for y in years:
            target_keys.add(f"{int(y):04d}")
    if months:
        for m in months:
            target_keys.add(parse_month_token(m))
    if not target_keys:
        return []
    chosen: List[Tuple[str, str, str]] = []
    if prefer_monthly:
        # 'YYYY' 만 지정한 경우 → 그 해의 모든 월별 파일을 우선 선택
        for k, c, u in index:
            if k in target_keys:
                chosen.append((k, c, u))
                continue
            if len(k) == 6 and k[:4] in target_keys:
                chosen.append((k, c, u))
    else:
        for k, c, u in index:
            if k in target_keys:
                chosen.append((k, c, u))
    return chosen


def download_with_progress(url: str, dest: Path) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    if dest.exists() and dest.stat().st_size > 0:
        print(f"[skip] {dest.name} (이미 있음 {dest.stat().st_size//1024} KB)")
        return
    print(f"[download] {url}")
    req = urllib.request.Request(url, headers={"User-Agent": "baduk-vlm-server/0.1"})
    with urllib.request.urlopen(req, timeout=120) as resp, open(dest, "wb") as f:
        total = int(resp.headers.get("Content-Length") or 0)
        read = 0
        chunk = 1024 * 64
        while True:
            buf = resp.read(chunk)
            if not buf:
                break
            f.write(buf)
            read += len(buf)
            if total:
                pct = read * 100 // total
                print(f"  {pct:3d}% ({read//1024}/{total//1024} KB)", end="\r", flush=True)
        print()


def extract_tar_gz(tar_path: Path, out_dir: Path) -> int:
    out_dir.mkdir(parents=True, exist_ok=True)
    n = 0
    # tar 내부에 디렉토리 구조가 잡혀 있어도, 모든 .sgf 를 out_dir 평면에 풉니다.
    with tarfile.open(tar_path, "r:gz") as tf:
        for m in tf:
            if not m.isfile():
                continue
            name = os.path.basename(m.name)
            if not name.lower().endswith(".sgf"):
                continue
            target = out_dir / name
            if target.exists() and target.stat().st_size > 0:
                continue
            f = tf.extractfile(m)
            if f is None:
                continue
            with open(target, "wb") as out_f:
                out_f.write(f.read())
            n += 1
    return n


def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(description="u-go.net KGS4d SGF 자동 다운로더")
    p.add_argument("--years", nargs="*", default=None, help="예: 2018 2017 (해당 해의 월별 모두)")
    p.add_argument("--months", nargs="*", default=None, help="예: 2018-01 2018-02")
    p.add_argument("--out-dir", default="sgf_data", help="압축 풀 폴더")
    p.add_argument("--archive-dir", default="sgf_data/_archives", help="원본 tar.gz 보관")
    p.add_argument("--list-only", action="store_true", help="다운로드 후보만 보고 끝.")
    p.add_argument("--max-archives", type=int, default=0, help="최대 N개만 받기 (0=무제한)")
    args = p.parse_args(argv)

    if not args.years and not args.months:
        p.error("--years 또는 --months 중 하나는 지정해야 합니다.")

    index = fetch_index()
    chosen = select_archives(index, args.years, args.months)
    if not chosen:
        print("선택된 아카이브가 없습니다. (인덱스에 해당 기간이 없거나 잘못 지정했습니다)")
        return 2

    print(f"총 {len(chosen)}개 아카이브 선택됨:")
    for k, c, u in chosen:
        print(f"  {k}  {c} games  {u}")
    if args.list_only:
        return 0

    if args.max_archives > 0:
        chosen = chosen[: args.max_archives]
        print(f"--max-archives {args.max_archives} → {len(chosen)} 개만 받습니다.")

    arch_dir = Path(args.archive_dir)
    sgf_dir = Path(args.out_dir)
    arch_dir.mkdir(parents=True, exist_ok=True)
    sgf_dir.mkdir(parents=True, exist_ok=True)

    total_extracted = 0
    for k, c, u in chosen:
        fname = os.path.basename(urllib.parse.urlparse(u).path)
        dest = arch_dir / fname
        try:
            download_with_progress(u, dest)
        except Exception as e:
            print(f"[fail] {u}: {e}")
            continue
        try:
            n = extract_tar_gz(dest, sgf_dir)
            total_extracted += n
            print(f"[extract] {fname} → {n} sgf files")
        except Exception as e:
            print(f"[extract fail] {dest}: {e}")
            continue

    summary = {
        "out_dir": str(sgf_dir.resolve()),
        "archives": [u for _, _, u in chosen],
        "extracted_files": total_extracted,
    }
    print(json.dumps(summary, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    sys.exit(main())