Spaces:
Running
Running
File size: 2,769 Bytes
5c9b605 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """A-share stock-code normalization helpers."""
from __future__ import annotations
from dataclasses import dataclass
@dataclass(frozen=True)
class NormalizedStockCode:
code: str
market: str
suffix: str
prefixed: str
display: str
class InvalidSymbolError(ValueError):
def __init__(self, symbol: str, message: str | None = None) -> None:
self.symbol = symbol
self.code = "invalid_symbol"
super().__init__(
message
or f"invalid_symbol: {symbol} is not a valid A-share stock code; "
"ETF/fund codes should use ETF or fund endpoints"
)
def normalize_stock_code(stock_code: str) -> NormalizedStockCode:
raw = str(stock_code or "").strip()
if not raw:
raise InvalidSymbolError(stock_code, "invalid_symbol: stock_code is required")
lower = raw.lower()
if lower.startswith(("sh", "sz", "bj")) and len(lower) >= 8:
market = lower[:2]
code = lower[2:8]
elif "." in raw:
code_part, suffix_part = raw.split(".", 1)
code = code_part.strip()
market = suffix_part.strip().lower()
else:
code = raw[:6]
inferred = _infer_a_share_market(code)
if inferred:
market = inferred
elif code.startswith(("8", "4", "920")):
market = "bj"
elif code.startswith(("510", "512", "513", "515", "588", "159")):
market = "sz" if code.startswith("159") else "sh"
elif code.startswith(("600", "601", "603", "605", "688")):
market = "sh"
else:
market = "sz"
if len(code) != 6 or not code.isdigit():
raise InvalidSymbolError(stock_code)
if market not in {"sh", "sz", "bj"}:
raise InvalidSymbolError(stock_code, f"invalid_symbol: unsupported market suffix for {stock_code}")
expected_market = _infer_a_share_market(code)
if expected_market is None:
raise InvalidSymbolError(stock_code)
if market != expected_market:
raise InvalidSymbolError(
stock_code,
f"invalid_symbol: {stock_code} does not match expected {expected_market.upper()} market for A-share stocks",
)
suffix = market.upper()
return NormalizedStockCode(
code=code,
market=market,
suffix=suffix,
prefixed=f"{market}{code}",
display=f"{code}.{suffix}",
)
def _infer_a_share_market(code: str) -> str | None:
if len(code) != 6 or not code.isdigit():
return None
if code.startswith(("600", "601", "603", "605", "688")):
return "sh"
if code.startswith(("000", "001", "002", "003", "300", "301")):
return "sz"
if code.startswith(("920", "8", "4")):
return "bj"
return None
|