Files
intl_news/extractor/pipeline.py
T
2026-07-18 16:13:52 +08:00

176 lines
5.2 KiB
Python

"""正文提取管道:批量处理 raw 目录下的文章"""
import json
import logging
from datetime import datetime
from pathlib import Path
from crawler.storage import load_index
from crawler.utils import get_news_day
from extractor.extractor import extract_article
from extractor.models import ProcessedArticle
logger = logging.getLogger(__name__)
def get_raw_data_sources(base_dir: str = "data/raw") -> list[str]:
"""扫描 data/raw/ 下所有源 ID
Args:
base_dir: raw 数据根目录
Returns:
源 ID 列表
"""
raw_path = Path(base_dir)
if not raw_path.exists():
return []
return sorted([
d.name for d in raw_path.iterdir()
if d.is_dir() and not d.name.startswith(".")
])
def process_source(
source_id: str,
date_str: str | None = None,
) -> list[ProcessedArticle]:
"""处理单个源的文章正文提取
Args:
source_id: 新闻源 ID
date_str: 日期 YYYYMMDD,默认当前新闻日
Returns:
ProcessedArticle 列表
"""
if date_str is None:
date_str = get_news_day()
logger.info("━━━ 正文提取 [%s] %s ━━━", source_id, date_str)
# 读取 raw index
articles = load_index(source_id, date_str)
if not articles:
logger.warning("[%s] %s 无待处理文章", source_id, date_str)
return []
# 输出目录
out_dir = Path(f"data/processed/{source_id}/{date_str}")
out_dir.mkdir(parents=True, exist_ok=True)
results: list[ProcessedArticle] = []
success = 0
no_content = 0
failed = 0
# 增量:跳过已处理的文章
skipped = 0
for article in articles[:]:
out_file = out_dir / f"{article.url_hash}.json"
if out_file.exists():
articles.remove(article)
skipped += 1
if skipped > 0:
logger.info("[%s] 增量跳过 %d 篇已处理,剩余 %d 篇",
source_id, skipped, len(articles))
for article in articles:
# 确保 html_path / md_path 是绝对路径(相对于 data/raw/)
html_path = str(Path("data/raw") / source_id / date_str / f"{article.url_hash}.html")
md_path = str(Path("data/raw") / source_id / date_str / f"{article.url_hash}.md")
# 如果 index 中已有路径,优先使用
if article.html_path:
html_path = article.html_path
if article.md_path:
md_path = article.md_path
processed = extract_article(
source_id=source_id,
source_name=article.source_name,
url=article.url,
url_hash=article.url_hash,
html_path=html_path,
md_path=md_path,
crawl_publish_time=article.publish_time,
)
# 写入 JSON 输出
out_file = out_dir / f"{article.url_hash}.json"
out_file.write_text(
processed.model_dump_json(indent=2, ensure_ascii=False),
encoding="utf-8",
)
results.append(processed)
if processed.status == "success":
success += 1
logger.debug("[%s] ✅ %s (%d words, %s)",
source_id, processed.title[:40], processed.word_count, processed.extractor)
elif processed.status == "no_content":
no_content += 1
logger.warning("[%s] ⚠️ 无正文: %s", source_id, processed.url[:80])
else:
failed += 1
logger.info("[%s] 完成: 成功 %d / 无内容 %d / 失败 %d",
source_id, success, no_content, failed)
# 写入处理索引
index_path = out_dir / "index.jsonl"
with open(index_path, "a", encoding="utf-8") as f:
for r in results:
if r.status == "success":
f.write(json.dumps({
"url_hash": r.url_hash,
"title": r.title,
"word_count": r.word_count,
"extractor": r.extractor,
"publish_time": r.publish_time,
}, ensure_ascii=False) + "\n")
return results
def process_all_sources(
source_filter: str | None = None,
date_str: str | None = None,
) -> dict:
"""处理所有源的文章正文提取
Args:
source_filter: 可选,只处理指定源
date_str: 日期,默认当前新闻日
Returns:
统计 dict
"""
if date_str is None:
date_str = get_news_day()
start_time = datetime.now()
if source_filter:
sources = [source_filter] if source_filter in get_raw_data_sources() else []
else:
sources = get_raw_data_sources()
logger.info("══════ 开始正文提取 %d 个源,日期: %s ══════", len(sources), date_str)
total = 0
for src in sources:
results = process_source(src, date_str)
total += len([r for r in results if r.status == "success"])
elapsed = (datetime.now() - start_time).total_seconds()
logger.info("══════ 提取完成: %d 篇文章,耗时 %.1f 秒 ══════", total, elapsed)
return {
"sources_processed": len(sources),
"total_articles": total,
"elapsed_sec": elapsed,
"date": date_str,
}