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,
tags TEXT,
metadata TEXT,
name TEXT,
industry TEXT,
created_at TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
PRIMARY KEY (trade_date, ts_code)
);

View File

@ -475,6 +475,7 @@ def _add_to_stock_pool(
) -> None:
"""将股票评分结果写入投资池。"""
broker = DataBroker()
trade_date = eval_date.strftime("%Y%m%d")
payload: List[tuple] = []
ranked_df = score_df.reset_index(drop=True)
@ -489,6 +490,11 @@ def _add_to_stock_pool(
},
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(
(
trade_date,
@ -498,6 +504,8 @@ def _add_to_stock_pool(
"factor_evaluation_top20",
tags,
metadata,
stock_name,
stock_industry,
)
)
@ -513,8 +521,10 @@ def _add_to_stock_pool(
status,
rationale,
tags,
metadata
) VALUES (?, ?, ?, ?, ?, ?, ?)
metadata,
name,
industry
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
payload,
)

View File

@ -45,7 +45,7 @@ def list_investment_pool(
"""Return investment candidates for the given trade date (latest if None)."""
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",
]
params: List[Any] = []

View File

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

View File

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