128 lines
4.3 KiB
Python
128 lines
4.3 KiB
Python
"""
|
||
报销全流程编排
|
||
|
||
将发票提取 → OCR 识别 → 浏览器填报串联为一条管道，
|
||
数据在内存中流转，同时生成 CSV / Markdown 中间产物。
|
||
"""
|
||
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
from . import get_logger
|
||
from .config import load_config
|
||
from .extractor import extract_invoices, save_csv as save_invoice_csv, save_markdown as save_invoice_md
|
||
from .ocr import enrich_with_ocr, _save_csv as save_ocr_csv, save_markdown_from_csv, _load_csv
|
||
|
||
log = get_logger("pipeline")
|
||
|
||
|
||
def run_pipeline(step: str = "all", username: str = None, password: str = None):
|
||
"""执行报销流程
|
||
|
||
Args:
|
||
step: all | invoice | ocr | submit
|
||
username: 覆盖 config.json 中的用户名
|
||
password: 覆盖 config.json 中的密码
|
||
"""
|
||
config = load_config()
|
||
if username:
|
||
config["username"] = username
|
||
if password:
|
||
config["password"] = password
|
||
|
||
# 工作目录（项目根目录）
|
||
project_dir = Path(__file__).parent.parent
|
||
|
||
# --------------------------------------------------
|
||
# Step 1: 发票提取
|
||
# --------------------------------------------------
|
||
invoices = None
|
||
|
||
if step in ("all", "invoice"):
|
||
log.info("=" * 60)
|
||
log.info("[1/3] 发票提取")
|
||
log.info("=" * 60)
|
||
|
||
invoices = extract_invoices(str(project_dir))
|
||
if not invoices:
|
||
log.error("未提取到任何发票数据")
|
||
return 1
|
||
|
||
save_invoice_csv(invoices, project_dir / "invoice_summary.csv")
|
||
save_invoice_md(invoices, project_dir / "invoice_summary.md")
|
||
|
||
if step == "invoice":
|
||
log.info("[1/3] 发票提取 完成")
|
||
return 0
|
||
|
||
# --------------------------------------------------
|
||
# Step 2: OCR 识别
|
||
# --------------------------------------------------
|
||
if step in ("all", "ocr"):
|
||
log.info("=" * 60)
|
||
log.info("[2/3] OCR 识别")
|
||
log.info("=" * 60)
|
||
|
||
csv_path = project_dir / "invoice_summary.csv"
|
||
|
||
if invoices is None:
|
||
rows = _load_csv(csv_path)
|
||
if rows is None:
|
||
return 1
|
||
else:
|
||
# 将 dict 列表转为 CSV 风格的 dict（对齐列名）
|
||
from .extractor import CSV_COLUMNS
|
||
rows = []
|
||
for idx, inv in enumerate(invoices, 1):
|
||
items = inv.get("_items", [])
|
||
first_item = items[0] if items else {}
|
||
rows.append({
|
||
"序号": str(idx),
|
||
"发票号码": inv.get("发票号码", ""),
|
||
"开票日期": inv.get("开票日期", ""),
|
||
"项目名称": first_item.get("项目名称", inv.get("项目名称", "")),
|
||
"规格型号": first_item.get("规格型号", inv.get("规格型号", "")),
|
||
"价税合计": inv.get("价税合计", ""),
|
||
"销售方名称": inv.get("销售方名称", ""),
|
||
"人员姓名": inv.get("人员姓名", ""),
|
||
"刷卡日期": inv.get("刷卡日期", ""),
|
||
"公务卡号": inv.get("公务卡号", ""),
|
||
"刷卡金额": inv.get("刷卡金额", ""),
|
||
"备注": inv.get("备注", ""),
|
||
"工号": inv.get("工号", ""),
|
||
})
|
||
|
||
rows = enrich_with_ocr(rows, str(project_dir))
|
||
save_ocr_csv(csv_path, rows)
|
||
save_markdown_from_csv(csv_path, rows)
|
||
invoices = rows
|
||
|
||
if step == "ocr":
|
||
log.info("[2/3] OCR 识别 完成")
|
||
return 0
|
||
|
||
# --------------------------------------------------
|
||
# Step 3: 浏览器填报
|
||
# --------------------------------------------------
|
||
if step in ("all", "submit"):
|
||
log.info("=" * 60)
|
||
log.info("[3/3] 报销提交")
|
||
log.info("=" * 60)
|
||
|
||
from .bot import load_invoice_data, run_bot
|
||
|
||
csv_path = project_dir / "invoice_summary.csv"
|
||
bot_invoices = load_invoice_data(str(csv_path), config)
|
||
run_bot(config, bot_invoices)
|
||
|
||
if step == "submit":
|
||
log.info("[3/3] 报销提交 完成")
|
||
return 0
|
||
|
||
# --------------------------------------------------
|
||
# 全流程完成
|
||
# --------------------------------------------------
|
||
log.info("=" * 60)
|
||
log.info("全流程执行完毕")
|
||
log.info("=" * 60)
|
||
return 0 |