Spaces:
Running
Running
| """A-share stock-code normalization helpers.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| 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 | |