Yuxuan Yan
95 lines
2.9 KiB
Python
95 lines
2.9 KiB
Python
"""Symbol conversion between internal A-share ids and JoinQuant ids."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
|
|
|
|
_INTERNAL_RE = re.compile(r"^(?P<prefix>sh|sz)(?P<code>\d{6})$", re.IGNORECASE)
|
|
_JOINQUANT_RE = re.compile(
|
|
r"^(?P<code>\d{6})\.(?P<exchange>XSHG|XSHE)$", re.IGNORECASE
|
|
)
|
|
_BARE_RE = re.compile(r"^\d{6}$")
|
|
|
|
|
|
def _validate_internal(prefix: str, code: str) -> None:
|
|
if prefix == "sh" and code.startswith("6"):
|
|
return
|
|
if prefix == "sz" and code.startswith(("0", "3")):
|
|
return
|
|
raise ValueError(f"Unsupported A-share symbol: {prefix}{code}")
|
|
|
|
|
|
def _validate_joinquant(code: str, exchange: str) -> None:
|
|
if exchange == "XSHG" and code.startswith("6"):
|
|
return
|
|
if exchange == "XSHE" and code.startswith(("0", "3")):
|
|
return
|
|
raise ValueError(f"Unsupported JoinQuant A-share symbol: {code}.{exchange}")
|
|
|
|
|
|
def to_joinquant_symbol(symbol_id: str) -> str:
|
|
"""Convert an internal symbol like ``sh600000`` to ``600000.XSHG``.
|
|
|
|
Args:
|
|
symbol_id: Internal A-share id with ``sh`` or ``sz`` prefix.
|
|
|
|
Returns:
|
|
JoinQuant security id.
|
|
|
|
Raises:
|
|
ValueError: If the symbol is malformed or outside supported A-share
|
|
Shanghai/Shenzhen equity prefixes.
|
|
"""
|
|
text = str(symbol_id).strip().lower()
|
|
match = _INTERNAL_RE.match(text)
|
|
if not match:
|
|
raise ValueError(f"Invalid internal symbol: {symbol_id!r}")
|
|
|
|
prefix = match.group("prefix").lower()
|
|
code = match.group("code")
|
|
_validate_internal(prefix, code)
|
|
exchange = "XSHG" if prefix == "sh" else "XSHE"
|
|
return f"{code}.{exchange}"
|
|
|
|
|
|
def from_joinquant_symbol(jq_symbol: str) -> str:
|
|
"""Convert a JoinQuant symbol like ``600000.XSHG`` to ``sh600000``."""
|
|
text = str(jq_symbol).strip().upper()
|
|
match = _JOINQUANT_RE.match(text)
|
|
if not match:
|
|
raise ValueError(f"Invalid JoinQuant symbol: {jq_symbol!r}")
|
|
|
|
code = match.group("code")
|
|
exchange = match.group("exchange").upper()
|
|
_validate_joinquant(code, exchange)
|
|
prefix = "sh" if exchange == "XSHG" else "sz"
|
|
return f"{prefix}{code}"
|
|
|
|
|
|
def normalize_symbol_pair(value: object) -> tuple[str, str]:
|
|
"""Return ``(symbol_id, jq_symbol)`` from any supported symbol spelling."""
|
|
text = str(value).strip()
|
|
if not text or text.lower() == "nan":
|
|
raise ValueError("Missing symbol")
|
|
|
|
if _INTERNAL_RE.match(text):
|
|
symbol_id = text.lower()
|
|
return symbol_id, to_joinquant_symbol(symbol_id)
|
|
|
|
if _JOINQUANT_RE.match(text):
|
|
jq_symbol = text.upper()
|
|
return from_joinquant_symbol(jq_symbol), jq_symbol
|
|
|
|
if _BARE_RE.match(text):
|
|
if text.startswith("6"):
|
|
symbol_id = f"sh{text}"
|
|
elif text.startswith(("0", "3")):
|
|
symbol_id = f"sz{text}"
|
|
else:
|
|
raise ValueError(f"Unsupported bare A-share code: {text}")
|
|
return symbol_id, to_joinquant_symbol(symbol_id)
|
|
|
|
raise ValueError(f"Unsupported symbol: {value!r}")
|
|
|