public-resources/company-deep/scripts/query_company.py

# -*- coding: utf-8 -*-
"""公司深度研报查询脚本 -- 供 Claude Skill 使用。

用法:
    python query_company.py <command> [args...]

子命令:
    lookup <keyword>              按关键词或股票代码查找公司
    list                          列出所有公司
    parts <stock_code>            列出该公司所有 Part 目录（编号+标题+字数）
    part <stock_code> <nums|all>  获取 Part 原始 Markdown
"""
from __future__ import annotations

import json
import sys
import os

sys.stdout.reconfigure(encoding="utf-8")
sys.stderr.reconfigure(encoding="utf-8")

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from db_config import DB_CONFIG

try:
    import pymysql
except ImportError:
    import subprocess
    print("pymysql 未安装，正在自动安装...", file=sys.stderr)
    subprocess.check_call([sys.executable, "-m", "pip", "install", "pymysql", "-q"])
    import pymysql


def _get_conn():
    return pymysql.connect(**DB_CONFIG, cursorclass=pymysql.cursors.DictCursor)


def cmd_lookup(keyword: str):
    """按关键词查找公司，模糊匹配 company_name 和 stock_code。"""
    conn = _get_conn()
    try:
        with conn.cursor() as cur:
            like = f"%{keyword}%"
            cur.execute(
                """SELECT stock_code, company_name, report_version,
                          total_parts, success_parts, total_characters,
                          generated_at
                   FROM company_reports
                   WHERE company_name LIKE %s OR stock_code LIKE %s
                   ORDER BY stock_code""",
                (like, like),
            )
            rows = cur.fetchall()
    finally:
        conn.close()

    if not rows:
        print(f"未找到匹配「{keyword}」的公司")
        return

    print(f"找到 {len(rows)} 家公司:\n")
    for r in rows:
        parts_info = f"{r['success_parts']}/{r['total_parts']} Part"
        chars = f"{r['total_characters']:,}字" if r["total_characters"] else ""
        print(f"  {r['stock_code']}  {r['company_name']}  {r['report_version']}  {parts_info}  {chars}")


def cmd_list():
    """列出所有公司。"""
    conn = _get_conn()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """SELECT stock_code, company_name, total_parts, success_parts,
                          total_characters, generated_at
                   FROM company_reports ORDER BY stock_code"""
            )
            rows = cur.fetchall()
    finally:
        conn.close()

    if not rows:
        print("暂无公司数据")
        return

    print(f"共 {len(rows)} 家公司:\n")
    for r in rows:
        parts_info = f"{r['success_parts']}/{r['total_parts']}"
        chars = f"{r['total_characters']:,}字" if r["total_characters"] else ""
        print(f"  {r['stock_code']}  {r['company_name']}  Parts={parts_info}  {chars}")


def cmd_parts(stock_code: str):
    """列出某公司所有 Part 的编号、标题、字数。"""
    conn = _get_conn()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """SELECT rrp.part_num, rrp.title, rrp.char_count
                   FROM report_raw_parts rrp
                   JOIN company_reports cr ON cr.id = rrp.report_id
                   WHERE cr.stock_code = %s
                   ORDER BY rrp.part_num""",
                (stock_code,),
            )
            rows = cur.fetchall()
    finally:
        conn.close()

    if not rows:
        print(f"未找到公司: {stock_code}")
        return

    print(f"公司 {stock_code} 的 Part 列表:\n")
    print(f"  {'Part':>6s}  {'标题':<30s}  {'字数':>8s}")
    print(f"  {'-'*6}  {'-'*30}  {'-'*8}")
    for r in rows:
        pn = r["part_num"]
        label = f"Part {pn}" if pn > 0 else "meta"
        title = r["title"] or ""
        print(f"  {label:>6s}  {title:<30s}  {r['char_count']:>8,d}")


def cmd_part(stock_code: str, part_nums: str):
    """获取 Part 原始 Markdown 文本。"""
    conn = _get_conn()
    try:
        with conn.cursor() as cur:
            if part_nums == "all":
                cur.execute(
                    """SELECT rrp.part_num, rrp.title, rrp.content, rrp.char_count
                       FROM report_raw_parts rrp
                       JOIN company_reports cr ON cr.id = rrp.report_id
                       WHERE cr.stock_code = %s
                       ORDER BY rrp.part_num""",
                    (stock_code,),
                )
            else:
                nums = [int(x.strip()) for x in part_nums.split(",")]
                placeholders = ",".join(["%s"] * len(nums))
                cur.execute(
                    f"""SELECT rrp.part_num, rrp.title, rrp.content, rrp.char_count
                        FROM report_raw_parts rrp
                        JOIN company_reports cr ON cr.id = rrp.report_id
                        WHERE cr.stock_code = %s AND rrp.part_num IN ({placeholders})
                        ORDER BY rrp.part_num""",
                    [stock_code] + nums,
                )
            rows = cur.fetchall()
    finally:
        conn.close()

    if not rows:
        print(f"未找到公司 {stock_code} 的 Part 数据")
        return

    for r in rows:
        pn = r["part_num"]
        label = f"Part {pn}: {r['title']}" if r["title"] else f"Part {pn}"
        if pn == 0:
            label = "meta (目录+Executive Summary)"
        print(f"{'=' * 60}")
        print(f"  {label}  ({r['char_count']:,} 字)")
        print(f"{'=' * 60}")
        print(r["content"])
        print()


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(0)

    cmd = sys.argv[1]

    if cmd == "lookup":
        if len(sys.argv) < 3:
            print("用法: query_company.py lookup <keyword>")
            sys.exit(1)
        cmd_lookup(sys.argv[2])

    elif cmd == "list":
        cmd_list()

    elif cmd == "parts":
        if len(sys.argv) < 3:
            print("用法: query_company.py parts <stock_code>")
            sys.exit(1)
        cmd_parts(sys.argv[2])

    elif cmd == "part":
        if len(sys.argv) < 4:
            print("用法: query_company.py part <stock_code> <0-15|all>")
            sys.exit(1)
        cmd_part(sys.argv[2], sys.argv[3])

    else:
        print(f"未知命令: {cmd}")
        print(__doc__)
        sys.exit(1)


if __name__ == "__main__":
    main()