Files
xiaobai-review/xiaobai-datahub/datahub/coverage.py
T
2026-09-02 22:16:05 +08:00

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