Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: multica-agent <github@multica.ai>
131 lines
3.9 KiB
Python
131 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Iterable
|
|
|
|
from datahub.db import HubDB
|
|
from datahub.timeutil import iter_yyyymmdd, yyyymmdd
|
|
|
|
MISSING_SAMPLE_LIMIT = 10
|
|
|
|
|
|
def coverage_payload(
|
|
*,
|
|
kind: str,
|
|
start: str,
|
|
end: str,
|
|
expected: Iterable[str],
|
|
available: Iterable[str],
|
|
extra: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
start = yyyymmdd(start)
|
|
end = yyyymmdd(end)
|
|
expected_list = sorted({yyyymmdd(item) for item in expected if item})
|
|
available_set = {yyyymmdd(item) for item in available if item}
|
|
missing = [item for item in expected_list if item not in available_set]
|
|
payload: dict[str, Any] = {
|
|
"kind": kind,
|
|
"complete": not missing,
|
|
"requested_from": start,
|
|
"requested_to": end,
|
|
"available_from": min(available_set) if available_set else None,
|
|
"available_to": max(available_set) if available_set else None,
|
|
"expected_count": len(expected_list),
|
|
"available_count": len(available_set),
|
|
"missing_count": len(missing),
|
|
"missing_sample": missing[:MISSING_SAMPLE_LIMIT],
|
|
}
|
|
if extra:
|
|
payload.update(extra)
|
|
return payload
|
|
|
|
|
|
def calendar_coverage(db: HubDB, start: str, end: str, exchange: str = "SSE") -> dict[str, Any]:
|
|
start = yyyymmdd(start)
|
|
end = yyyymmdd(end)
|
|
expected = list(iter_yyyymmdd(start, end))
|
|
rows = db.fetchall(
|
|
"SELECT cal_date FROM trade_calendar WHERE exchange = ? AND cal_date >= ? AND cal_date <= ?",
|
|
(exchange, start, end),
|
|
)
|
|
return coverage_payload(
|
|
kind="calendar",
|
|
start=start,
|
|
end=end,
|
|
expected=expected,
|
|
available=(row["cal_date"] for row in rows),
|
|
extra={"exchange": exchange},
|
|
)
|
|
|
|
|
|
def published_range_coverage(
|
|
db: HubDB,
|
|
dataset: str,
|
|
start: str,
|
|
end: str,
|
|
ts_code: str = "",
|
|
table: str = "",
|
|
) -> dict[str, Any]:
|
|
start = yyyymmdd(start)
|
|
end = yyyymmdd(end)
|
|
calendar = calendar_coverage(db, start, end)
|
|
open_rows = db.fetchall(
|
|
"""
|
|
SELECT cal_date FROM trade_calendar
|
|
WHERE exchange = 'SSE' AND is_open = 1 AND cal_date >= ? AND cal_date <= ?
|
|
ORDER BY cal_date
|
|
""",
|
|
(start, end),
|
|
)
|
|
expected_open = [row["cal_date"] for row in open_rows]
|
|
pubs = db.fetchall(
|
|
"""
|
|
SELECT trade_date, active_batch FROM publications
|
|
WHERE dataset = ? AND trade_date >= ? AND trade_date <= ?
|
|
ORDER BY trade_date
|
|
""",
|
|
(dataset, start, end),
|
|
)
|
|
published_dates = [row["trade_date"] for row in pubs]
|
|
available = list(published_dates)
|
|
extra: dict[str, Any] = {
|
|
"dataset": dataset,
|
|
"calendar_complete": calendar["complete"],
|
|
"calendar_missing_count": calendar["missing_count"],
|
|
}
|
|
if ts_code and table and pubs:
|
|
present_code: list[str] = []
|
|
for pub in pubs:
|
|
hit = db.fetchone(
|
|
f"SELECT 1 AS ok FROM {table} WHERE trade_date = ? AND batch_id = ? AND ts_code = ? LIMIT 1",
|
|
(pub["trade_date"], pub["active_batch"], ts_code),
|
|
)
|
|
if hit:
|
|
present_code.append(pub["trade_date"])
|
|
available = present_code
|
|
extra["code"] = ts_code
|
|
payload = coverage_payload(
|
|
kind="published_range",
|
|
start=start,
|
|
end=end,
|
|
expected=expected_open,
|
|
available=available,
|
|
extra=extra,
|
|
)
|
|
if not calendar["complete"]:
|
|
payload["complete"] = False
|
|
payload["calendar_missing_sample"] = calendar["missing_sample"]
|
|
return payload
|
|
|
|
|
|
def point_coverage(trade_date: str, dataset: str = "") -> dict[str, Any]:
|
|
day = yyyymmdd(trade_date)
|
|
payload = coverage_payload(
|
|
kind="point",
|
|
start=day,
|
|
end=day,
|
|
expected=[day],
|
|
available=[day],
|
|
extra={"dataset": dataset} if dataset else None,
|
|
)
|
|
return payload
|