add stock name and industry fields to investment pool

This commit is contained in:
sam 2025-10-11 19:16:42 +08:00
parent efa3a77787
commit 6810712232
5 changed files with 24 additions and 8 deletions

View File

@ -474,6 +474,8 @@ SCHEMA_STATEMENTS: Iterable[str] = (
rationale TEXT, rationale TEXT,
tags TEXT, tags TEXT,
metadata TEXT, metadata TEXT,
name TEXT,
industry TEXT,
created_at TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), created_at TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
PRIMARY KEY (trade_date, ts_code) PRIMARY KEY (trade_date, ts_code)
); );

View File

@ -475,6 +475,7 @@ def _add_to_stock_pool(
) -> None: ) -> None:
"""将股票评分结果写入投资池。""" """将股票评分结果写入投资池。"""
broker = DataBroker()
trade_date = eval_date.strftime("%Y%m%d") trade_date = eval_date.strftime("%Y%m%d")
payload: List[tuple] = [] payload: List[tuple] = []
ranked_df = score_df.reset_index(drop=True) ranked_df = score_df.reset_index(drop=True)
@ -489,6 +490,11 @@ def _add_to_stock_pool(
}, },
ensure_ascii=False, ensure_ascii=False,
) )
# 获取股票基本信息
stock_info = broker.get_stock_info(row["股票代码"], trade_date)
stock_name = stock_info.get("name", "") if stock_info else ""
stock_industry = stock_info.get("industry", "") if stock_info else ""
payload.append( payload.append(
( (
trade_date, trade_date,
@ -498,6 +504,8 @@ def _add_to_stock_pool(
"factor_evaluation_top20", "factor_evaluation_top20",
tags, tags,
metadata, metadata,
stock_name,
stock_industry,
) )
) )
@ -513,8 +521,10 @@ def _add_to_stock_pool(
status, status,
rationale, rationale,
tags, tags,
metadata metadata,
) VALUES (?, ?, ?, ?, ?, ?, ?) name,
industry
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
""", """,
payload, payload,
) )

View File

@ -45,7 +45,7 @@ def list_investment_pool(
"""Return investment candidates for the given trade date (latest if None).""" """Return investment candidates for the given trade date (latest if None)."""
query = [ query = [
"SELECT trade_date, ts_code, score, status, rationale, tags, metadata", "SELECT trade_date, ts_code, score, status, rationale, tags, metadata, name, industry",
"FROM investment_pool", "FROM investment_pool",
] ]
params: List[Any] = [] params: List[Any] = []

View File

@ -76,6 +76,8 @@ SCHEMA_STATEMENTS = [
rationale TEXT, rationale TEXT,
tags TEXT, -- JSON array tags TEXT, -- JSON array
metadata TEXT, -- JSON object metadata TEXT, -- JSON object
name TEXT,
industry TEXT,
PRIMARY KEY (trade_date, ts_code) PRIMARY KEY (trade_date, ts_code)
); );
""", """,

View File

@ -133,19 +133,21 @@ def test_list_investment_pool_orders_without_nulls(tmp_path):
rationale TEXT, rationale TEXT,
tags TEXT, tags TEXT,
metadata TEXT, metadata TEXT,
name TEXT,
industry TEXT,
PRIMARY KEY (trade_date, ts_code) PRIMARY KEY (trade_date, ts_code)
) )
""" """
) )
conn.executemany( conn.executemany(
""" """
INSERT INTO investment_pool (trade_date, ts_code, score, status, rationale, tags, metadata) INSERT INTO investment_pool (trade_date, ts_code, score, status, rationale, tags, metadata, name, industry)
VALUES (?, ?, ?, ?, ?, ?, ?) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
""", """,
[ [
("2024-01-01", "AAA", 0.8, "buy", "", None, None), ("2024-01-01", "AAA", 0.8, "buy", "", None, None, "Company A", "Technology"),
("2024-01-01", "BBB", None, "hold", "", None, None), ("2024-01-01", "BBB", None, "hold", "", None, None, "Company B", "Finance"),
("2024-01-01", "CCC", 0.9, "buy", "", None, None), ("2024-01-01", "CCC", 0.9, "buy", "", None, None, "Company C", "Healthcare"),
], ],
) )