feat0308 新增compyny-deep技能
This commit is contained in:
204
company-deep/scripts/query_company.py
Normal file
204
company-deep/scripts/query_company.py
Normal file
@@ -0,0 +1,204 @@
|
||||
# -*- 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()
|
||||
Reference in New Issue
Block a user