from __future__ import annotations import argparse import time from pathlib import Path import requests from auth_login import ( DEFAULT_YAOEX_AUTH_PATH, DEFAULT_YBM_COOKIE_PATH, PROJECT_ROOT, BrowserRuntime, load_cookie_file, load_dotenv, load_yaoex_auth_file, login_yaoex, login_ybm, save_cookie_file, save_yaoex_auth, ) from db import RepoUnavailable from job import run_batch from mall_config import YAOEX_MALL_NAME, YBM_MALL_NAME from repo import OracleSelectRepo from rules import Candidate from search_yaoex import YaoexAuth from search_yaoex import search_candidates as search_yaoex_candidates from search_ybm import search_candidates as search_ybm_candidates def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser( description="选品班次:捞表1 deal_flag=0,按条形码搜平台,结果写入表2。" ) parser.add_argument("--limit", type=int, default=None, help="本班最多领取 N 条") parser.add_argument( "--platform", choices=["all", "yaoex", "ybm"], default="all", help="all=两平台都跑", ) parser.add_argument("--secrets", type=Path, default=None, help="secrets.json 路径") parser.add_argument( "--yaoex-auth", type=Path, default=None, help="已登录的 1号药城认证 JSON(token/user_id/buyer_code),跳过浏览器", ) parser.add_argument( "--ybm-cookies", type=Path, default=None, help="已登录的药帮忙 cookie JSON,跳过浏览器", ) parser.add_argument( "--keep-browser", action="store_true", help="登录后不关 Chrome,搜完再留一段时间方便查看", ) parser.add_argument( "--hold-seconds", type=int, default=180, help="--keep-browser 时窗口再留多少秒,默认 180", ) parser.add_argument( "--fresh-login", action="store_true", help="忽略本地登录缓存,强制重新打开浏览器", ) return parser.parse_args(argv) def _platforms(choice: str) -> list[str]: if choice == "yaoex": return [YAOEX_MALL_NAME] if choice == "ybm": return [YBM_MALL_NAME] return [YAOEX_MALL_NAME, YBM_MALL_NAME] def main(argv: list[str] | None = None, repo=None) -> None: args = parse_args(argv) load_dotenv(PROJECT_ROOT / ".env") flow_start = time.perf_counter() if repo is None: try: repo = OracleSelectRepo.from_env() except RepoUnavailable as exc: print(f"数据库不可用:{exc}") raise SystemExit(2) from exc yaoex_auth: YaoexAuth | None = None ybm_cookies: dict[str, str] | None = None login_errors: dict[str, str] = {} runtime = BrowserRuntime(keep=args.keep_browser, hold_seconds=args.hold_seconds) platforms = _platforms(args.platform) try: if YAOEX_MALL_NAME in platforms: try: cached = args.yaoex_auth or ( None if args.fresh_login else DEFAULT_YAOEX_AUTH_PATH ) if cached and cached.exists(): yaoex_auth = load_yaoex_auth_file(cached) print(f"1号药城使用缓存登录:{cached}") else: yaoex_auth = login_yaoex(args.secrets, runtime=runtime) save_yaoex_auth(yaoex_auth) print(f"1号药城登录已缓存:{DEFAULT_YAOEX_AUTH_PATH}") print("1号药城认证就绪。") except SystemExit: raise except Exception as exc: login_errors[YAOEX_MALL_NAME] = str(exc) print(f"1号药城登录失败,本平台中止:{exc}") if YBM_MALL_NAME in platforms: try: cached = args.ybm_cookies or ( None if args.fresh_login else DEFAULT_YBM_COOKIE_PATH ) if cached and cached.exists(): ybm_cookies = load_cookie_file(cached) print(f"药帮忙使用缓存登录:{cached}") else: ybm_cookies = login_ybm(args.secrets, runtime=runtime) save_cookie_file(ybm_cookies) print(f"药帮忙登录已缓存:{DEFAULT_YBM_COOKIE_PATH}") print("药帮忙认证就绪。") except SystemExit: raise except Exception as exc: login_errors[YBM_MALL_NAME] = str(exc) print(f"药帮忙登录失败,本平台中止:{exc}") yaoex_session = requests.Session() ybm_session = requests.Session() yaoex_cache: dict[str, list[Candidate]] = {} ybm_cache: dict[str, list[Candidate]] = {} if ybm_cookies: for name, value in ybm_cookies.items(): ybm_session.cookies.set(name, value) searchers = {} if yaoex_auth is not None: def search_yaoex(query_key: str) -> list[Candidate]: assert yaoex_auth is not None return search_yaoex_candidates( yaoex_auth, query_key, session=yaoex_session, cache=yaoex_cache ) searchers[YAOEX_MALL_NAME] = search_yaoex if ybm_cookies is not None: def search_ybm(query_key: str) -> list[Candidate]: return search_ybm_candidates( ybm_cookies, query_key, session=ybm_session, cache=ybm_cache ) searchers[YBM_MALL_NAME] = search_ybm try: summary = run_batch( repo, platforms=platforms, searchers=searchers, login_errors=login_errors, limit=args.limit, ) finally: yaoex_session.close() ybm_session.close() elapsed = time.perf_counter() - flow_start print( f"领取={summary.claimed} ADDED={summary.added} " f"SKIPPED={summary.skipped} FAILED={summary.failed} " f"用时 {elapsed:.2f} 秒" ) print("已写入 rpa_cart_settlement_goods_data,push_status=0,未推订单中心,未加购。") finally: runtime.close() if __name__ == "__main__": main()