kwcli 0.1.0__py3-none-any.whl

kiwoom/core/ws_client.py

@@ -0,0 +1,262 @@
+import asyncio
+import json
+import logging
+from collections.abc import AsyncIterator, Callable
+from typing import Any
+
+from kiwoom.core.auth import KiwoomAuth, get_ws_base_url
+from kiwoom.core.errors import AUTH_RETRY_RETURN_CODES, normalize_return_code
+
+logger = logging.getLogger(__name__)
+
+try:
+    from websockets.asyncio.client import connect as _ws_connect
+except ImportError:  # pragma: no cover - compatibility fallback
+    try:
+        from websockets.client import connect as _ws_connect
+    except ImportError:  # pragma: no cover - dependency installed later in the flow
+        _ws_connect = None
+
+
+type ConnectFactory = Callable[..., Any]
+
+
+class WebSocketLoginError(RuntimeError):
+    def __init__(self, return_code: int, return_msg: str) -> None:
+        super().__init__(f"websocket login failed: {return_msg}")
+        self.return_code = return_code
+        self.return_msg = return_msg
+
+
+class KiwoomWebSocketClient:
+    def __init__(
+        self,
+        auth: KiwoomAuth,
+        *,
+        timeout_seconds: int = 30,
+        connect_factory: ConnectFactory | None = None,
+    ) -> None:
+        self.auth = auth
+        self.timeout_seconds = timeout_seconds
+        self._connect_factory = connect_factory or _default_connect
+        self._websocket: Any | None = None
+
+    @property
+    def is_connected(self) -> bool:
+        """Whether a live WebSocket connection is currently open."""
+        return self._websocket is not None
+
+    async def connect(self, *, api_url: str, retry_on_auth_failure: bool = True) -> None:
+        if self._websocket is not None:
+            await self.close()
+
+        path = _parse_api_url(api_url)
+        uri = f"{get_ws_base_url(self.auth.mode)}{path}"
+        try:
+            await self._connect_and_login(uri)
+        except WebSocketLoginError as exc:
+            if not retry_on_auth_failure or not _is_auth_retry_login_error(exc):
+                raise RuntimeError(f"websocket login failed: {exc.return_msg}") from exc
+            logger.info("websocket auth failed; refreshing token and retrying login")
+            await self.close()
+            self.auth.recover_from_auth_failure()
+            try:
+                await self._connect_and_login(uri)
+            except WebSocketLoginError as retry_exc:
+                raise RuntimeError(f"websocket login failed: {retry_exc.return_msg}") from retry_exc
+
+    async def request_once(
+        self,
+        *,
+        api_url: str,
+        body: dict[str, Any],
+    ) -> Any:
+        await self.connect(api_url=api_url)
+        try:
+            logger.info("websocket sent request")
+            await self.send(body)
+            try:
+                message = await asyncio.wait_for(self._receive_message(), timeout=self.timeout_seconds)
+            except TimeoutError as exc:
+                raise TimeoutError(f"websocket response timed out after {self.timeout_seconds}s") from exc
+            logger.info("websocket received response")
+            return message
+        finally:
+            await self.close()
+
+    async def subscribe(
+        self,
+        *,
+        api_url: str,
+        body: dict[str, Any],
+    ) -> None:
+        await self.connect(api_url=api_url)
+        logger.info("websocket sent subscribe request")
+        await self.send(body)
+
+    async def send(self, payload: Any) -> None:
+        """Send a packet over the current WebSocket connection."""
+        await self._send_packet(payload)
+
+    async def recv(self) -> Any:
+        """Receive the next non-ping, non-login message over the current connection."""
+        return await self._receive_message()
+
+    async def iter_messages(self) -> AsyncIterator[Any]:
+        if self._websocket is None:
+            raise RuntimeError("websocket is not connected")
+
+        while True:
+            try:
+                yield await self._receive_message()
+            except Exception as exc:  # pragma: no cover - exercised with real websocket client
+                if _is_connection_closed_exception(exc):
+                    logger.info("websocket connection closed by remote")
+                    break
+                raise
+
+    async def close(self) -> None:
+        if self._websocket is None:
+            return
+
+        await self._websocket.close()
+        self._websocket = None
+        logger.info("websocket closed")
+
+    async def _connect_and_login(self, uri: str) -> None:
+        logger.info("websocket connecting: %s", uri)
+        self._websocket = await self._connect_factory(
+            uri,
+            self.timeout_seconds,
+        )
+        logger.info("websocket connected")
+        await self._login()
+
+    async def _receive_message(self) -> Any:
+        while True:
+            parsed = await self._receive_non_ping_message()
+            if _is_login_message(parsed):
+                return_code = normalize_return_code(parsed.get("return_code")) or 0
+                if return_code != 0:
+                    raise WebSocketLoginError(
+                        return_code,
+                        str(parsed.get("return_msg", "unknown error")),
+                    )
+                logger.info("websocket login acknowledged")
+                continue
+
+            logger.info("websocket received message")
+            return parsed
+
+    async def _login(self) -> None:
+        if self._websocket is None:
+            raise RuntimeError("websocket is not connected")
+
+        login_packet = {
+            "trnm": "LOGIN",
+            "token": self.auth.get_access_token(),
+        }
+        logger.info("websocket sending LOGIN packet")
+        await self._send_packet(login_packet)
+        await self._await_login_ack()
+
+    async def _send_packet(self, payload: Any) -> None:
+        if self._websocket is None:
+            raise RuntimeError("websocket is not connected")
+        await self._websocket.send(_serialize_message(payload))
+
+    async def _await_login_ack(self) -> None:
+        while True:
+            parsed = await self._receive_non_ping_message()
+            if not _is_login_message(parsed):
+                raise RuntimeError("websocket login ack was not received before other messages")
+            return_code = normalize_return_code(parsed.get("return_code")) or 0
+            if return_code != 0:
+                raise WebSocketLoginError(
+                    return_code,
+                    str(parsed.get("return_msg", "unknown error")),
+                )
+            logger.info("websocket login acknowledged")
+            return
+
+    async def _receive_non_ping_message(self) -> Any:
+        if self._websocket is None:
+            raise RuntimeError("websocket is not connected")
+
+        while True:
+            message = await self._websocket.recv()
+            parsed = _parse_message(message)
+            if not _is_ping_message(parsed):
+                return parsed
+            logger.info("websocket received ping; echoing back")
+            await self._send_packet(parsed)
+            logger.info("websocket sent ping echo")
+
+
+def _is_auth_retry_login_error(exc: WebSocketLoginError) -> bool:
+    if exc.return_code in AUTH_RETRY_RETURN_CODES:
+        return True
+    return any(f"CODE={code}" in exc.return_msg for code in AUTH_RETRY_RETURN_CODES)
+
+
+async def _default_connect(uri: str, timeout_seconds: int) -> Any:
+    if _ws_connect is None:
+        raise RuntimeError("websockets dependency is not installed")
+    return await _ws_connect(
+        uri,
+        open_timeout=timeout_seconds,
+        ping_interval=None,
+    )
+
+
+def _parse_api_url(api_url: str) -> str:
+    normalized = api_url.strip()
+    if not normalized.startswith("/"):
+        raise ValueError("api_url must be an API path starting with '/'")
+    if normalized.endswith("/websocket"):
+        return normalized
+
+    path, _, api_id = normalized.rpartition("/")
+    if not path or not api_id:
+        raise ValueError("api_url must be an API path starting with '/'")
+
+    return path
+
+
+def _parse_message(message: Any) -> Any:
+    if isinstance(message, bytes):
+        message = message.decode("utf-8")
+
+    if isinstance(message, str):
+        stripped = message.strip()
+        if not stripped:
+            return message
+        try:
+            return json.loads(stripped)
+        except json.JSONDecodeError:
+            return message
+
+    return message
+
+
+def _is_login_message(message: Any) -> bool:
+    return isinstance(message, dict) and str(message.get("trnm", "")).upper() == "LOGIN"
+
+
+def _is_ping_message(message: Any) -> bool:
+    if isinstance(message, str):
+        return message.strip().upper() == "PING"
+    if isinstance(message, dict):
+        return str(message.get("trnm", "")).upper() == "PING"
+    return False
+
+
+def _serialize_message(message: Any) -> str:
+    if isinstance(message, str):
+        return message
+    return json.dumps(message, ensure_ascii=False)
+
+
+def _is_connection_closed_exception(exc: Exception) -> bool:
+    name = exc.__class__.__name__
+    return "ConnectionClosed" in name or name.endswith("ClosedOK") or name.endswith("ClosedError")

kiwoom/realtime/__init__.py

@@ -0,0 +1,30 @@
+from .decoders import decode_realtime_message, decode_realtime_message_named, decode_values
+from .events import (
+    is_terminal_event,
+    normalize_message_events,
+    resolve_topic,
+    topic_fanout,
+)
+from .packets import build_reg_packet, build_remove_packet
+from .stream import (
+    AsyncPubSub,
+    collect_realtime,
+    event_to_dataframe,
+    run_pubsub,
+)
+
+__all__ = [
+    "AsyncPubSub",
+    "build_reg_packet",
+    "build_remove_packet",
+    "collect_realtime",
+    "decode_realtime_message",
+    "decode_realtime_message_named",
+    "decode_values",
+    "event_to_dataframe",
+    "is_terminal_event",
+    "normalize_message_events",
+    "resolve_topic",
+    "run_pubsub",
+    "topic_fanout",
+]

kiwoom/realtime/decoders.py

@@ -0,0 +1,171 @@
+"""Decoder for Kiwoom WebSocket REAL messages.
+
+Translates numeric FID keys in ``values`` dicts to Korean field names. The
+``{fid: 한글명}`` mapping is derived from the Kiwoom WebSocket spec
+(``kiwoom_api_spec.json``): each realtime type's REAL ``values`` fields are the
+depth-2 rows nested under ``values`` in the response body. The spec is the single
+source of truth and names are returned verbatim; unknown FID keys pass through
+unchanged so no data is lost.
+
+Typical REAL message shape::
+
+    {
+        "trnm": "REAL",
+        "data": [
+            {
+                "type": "0B",
+                "name": "주식체결",
+                "item": "005930",
+                "values": {"10": "-82000", "15": "+12345", ...}
+            }
+        ]
+    }
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Any
+
+from kiwoom.realtime.schemas import get_realtime_schema
+
+
+@lru_cache(maxsize=None)
+def _fid_map(real_type: str) -> dict[str, str]:
+    """Extract ``{fid: 한글명}`` from the spec's REAL ``values`` rows (depth 2)."""
+    from kiwoom.specs import get_api_spec
+
+    try:
+        payload = get_api_spec(real_type)
+    except ValueError:
+        return {}
+
+    fid_map: dict[str, str] = {}
+    in_values = False
+    for row in payload.get("response", {}).get("body", []):
+        element = str(row.get("element", "")).strip()
+        depth = row.get("depth")
+        if element == "values" and depth == 1:
+            in_values = True
+            continue
+        if not in_values:
+            continue
+        if depth == 2 and element:
+            fid_map[element] = str(row.get("한글명", "")).strip()
+        elif isinstance(depth, int) and depth <= 1:
+            in_values = False
+    return fid_map
+
+
+def fid_display_map(real_type: str) -> dict[str, str]:
+    """Return the spec-defined FID -> display name mapping for a realtime type."""
+    return dict(_fid_map(real_type))
+
+
+def decode_values(real_type: str, values: dict[str, str]) -> dict[str, str]:
+    """Translate FID-keyed *values* dict to Korean field names.
+
+    Unknown FID keys are passed through unchanged so no data is lost.
+
+    Args:
+        real_type: REAL type ID, e.g. ``"0B"``.
+        values: Raw ``{"fid": "value"}`` dict from the server message.
+
+    Returns:
+        ``{"field_name": "value"}`` dict with human-readable keys.
+    """
+    fid_map = fid_display_map(real_type)
+    return {fid_map.get(k, k): v for k, v in values.items()}
+
+
+def decode_realtime_message(message: dict[str, Any]) -> dict[str, Any]:
+    """Decode FID keys in all ``data[*].values`` entries of a REAL message.
+
+    Non-REAL messages (e.g. ACK, LOGIN) are returned unchanged.
+
+    Args:
+        message: Parsed WebSocket message dict.
+
+    Returns:
+        Copy of *message* with ``values`` dicts translated to field names.
+    """
+    if not isinstance(message, dict):
+        return message
+
+    if str(message.get("trnm", "")).upper() != "REAL":
+        return message
+
+    data = message.get("data")
+    if not isinstance(data, list):
+        return message
+
+    decoded_data = []
+    for entry in data:
+        if not isinstance(entry, dict):
+            decoded_data.append(entry)
+            continue
+
+        real_type = entry.get("type", "")
+        values = entry.get("values")
+
+        if isinstance(values, dict) and real_type:
+            decoded_data.append({**entry, "values": decode_values(real_type, values)})
+        else:
+            decoded_data.append(entry)
+
+    return {**message, "data": decoded_data}
+
+
+def decode_realtime_message_named(message: dict[str, Any]) -> dict[str, Any]:
+    """Convert REAL message entries into named data objects.
+
+    Only ``trnm == "REAL"`` messages are decoded. Control frames such as ``REG``
+    and ``SYSTEM`` are returned unchanged. Values keep their original string
+    representation; this function changes keys/shape only and preserves unknown
+    FIDs under ``unknown``.
+    """
+    if not isinstance(message, dict):
+        return message
+
+    if str(message.get("trnm", "")).upper() != "REAL":
+        return message
+
+    data = message.get("data")
+    if not isinstance(data, list):
+        return message
+
+    named_data = []
+    for entry in data:
+        if not isinstance(entry, dict):
+            named_data.append(entry)
+            continue
+
+        real_type = str(entry.get("type", "")).strip()
+        values = entry.get("values")
+        schema = get_realtime_schema(real_type)
+        if not isinstance(values, dict) or schema is None:
+            named_data.append(entry)
+            continue
+
+        fields = schema.get("fields", {})
+        named_values: dict[str, Any] = {}
+        unknown: dict[str, Any] = {}
+        for fid, value in values.items():
+            field = fields.get(str(fid))
+            if field is None:
+                unknown[str(fid)] = value
+                continue
+            named_values[str(field["name"])] = value
+
+        named_entry: dict[str, Any] = {
+            "event": schema.get("event", real_type),
+            "type": real_type,
+            "name": entry.get("name") or schema.get("name", ""),
+            "code": entry.get("item", ""),
+            "data": named_values,
+        }
+        if unknown:
+            named_entry["unknown"] = unknown
+        named_data.append(named_entry)
+
+    return {**message, "data": named_data}

kiwoom/realtime/events.py

@@ -0,0 +1,98 @@
+"""Generic helpers for turning Kiwoom WebSocket messages into pub/sub events.
+
+These utilities are intentionally API-agnostic: they only depend on the shape of
+Kiwoom WebSocket frames (``trnm`` / ``data`` / ``values``), not on any specific
+realtime type. Examples keep their own field label table (``COLUMNS``) inline and
+pass it in, so a single example still shows which fields it produces while the
+routing/normalization plumbing lives here.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def normalize_message_events(
+    message: Any,
+    columns: dict[str, str] | None = None,
+) -> list[dict[str, Any]]:
+    """Convert a received WebSocket message into a list of consumer events.
+
+    A ``REAL`` frame is expanded into one event per ``data`` entry, with FID keys
+    in ``values`` translated through *columns* (unknown keys pass through). Any
+    other frame (REG/SYSTEM/CNSR*/...) is returned unchanged as a single event.
+
+    Args:
+        message: Parsed WebSocket message.
+        columns: Optional ``{fid: label}`` map applied to ``values``.
+
+    Returns:
+        A list of event dicts ready to publish.
+    """
+    if not isinstance(message, dict):
+        return [{"trnm": "RAW", "message": message}]
+
+    trnm = str(message.get("trnm", "")).upper()
+    if trnm != "REAL":
+        return [{**message, "trnm": trnm or message.get("trnm", "")}]
+
+    column_map = columns or {}
+    events: list[dict[str, Any]] = []
+    data = message.get("data", [])
+    if not isinstance(data, list):
+        return [{"trnm": "REAL", "message": message}]
+
+    for entry in data:
+        if not isinstance(entry, dict):
+            continue
+        raw_values = entry.get("values")
+        raw_values = raw_values if isinstance(raw_values, dict) else {}
+        values = {column_map.get(key, key): value for key, value in raw_values.items()}
+        event: dict[str, Any] = {
+            "trnm": "REAL",
+            "item": entry.get("item", ""),
+            "type": str(entry.get("type", "")),
+        }
+        if "name" in entry:
+            event["name"] = entry.get("name", "")
+        event.update(values)
+        events.append(event)
+    return events
+
+
+def resolve_topic(event: dict[str, Any]) -> str:
+    """Decide the representative topic for a normalized *event*."""
+    if not isinstance(event, dict):
+        return "kiwoom.raw"
+
+    trnm = str(event.get("trnm", "")).upper()
+    if trnm == "REAL":
+        realtime_type = str(event.get("type", "")).strip()
+        if realtime_type:
+            return f"kiwoom.realtime.{realtime_type}"
+        return "kiwoom.realtime"
+    if trnm == "REG":
+        return "kiwoom.system.reg"
+    if trnm == "SYSTEM":
+        return "kiwoom.system"
+    if trnm:
+        return f"kiwoom.system.{trnm.lower()}"
+    return "kiwoom.raw"
+
+
+def topic_fanout(topic: str) -> list[str]:
+    """Expand a representative topic to the shared topics it also publishes to."""
+    topics = [topic]
+    if topic.startswith("kiwoom.realtime."):
+        topics.append("kiwoom.realtime")
+    if topic.startswith("kiwoom.system."):
+        topics.append("kiwoom.system")
+    topics.append("kiwoom.all")
+    return list(dict.fromkeys(topics))
+
+
+def is_terminal_event(event: dict[str, Any]) -> bool:
+    """Whether *event* should stop the receive loop (SYSTEM or error return_code)."""
+    trnm = str(event.get("trnm", "")).upper()
+    return_code = event.get("return_code")
+    return trnm == "SYSTEM" or return_code not in (None, 0, "0")

kiwoom/realtime/packets.py

@@ -0,0 +1,65 @@
+"""Builders for Kiwoom realtime WebSocket control packets.
+
+These helpers only encode the Kiwoom realtime request contract (``REG`` to
+register, ``REMOVE`` to unregister). They are API-agnostic: callers pass the
+items and realtime types they want, so a single builder serves every realtime
+example without per-API logic.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def build_reg_packet(
+    items: list[str],
+    types: list[str],
+    *,
+    group_no: str = "1",
+    refresh: str = "1",
+) -> dict[str, Any]:
+    """Build a Kiwoom realtime registration (``REG``) packet.
+
+    Args:
+        items: Symbols or elements to register (e.g. ``["005930"]``).
+        types: Realtime type codes to register (e.g. ``["0B"]``).
+        group_no: Group number.
+        refresh: ``"1"`` keeps existing registrations, ``"0"`` replaces them.
+
+    Returns:
+        A packet dict ready to send over the realtime WebSocket.
+    """
+    if not types:
+        raise ValueError("types is required.")
+    return {
+        "trnm": "REG",
+        "grp_no": group_no,
+        "refresh": refresh,
+        "data": [
+            {
+                "item": items,
+                "type": types,
+            }
+        ],
+    }
+
+
+def build_remove_packet(
+    items: list[str],
+    types: list[str],
+    *,
+    group_no: str = "1",
+) -> dict[str, Any]:
+    """Build a Kiwoom realtime unregistration (``REMOVE``) packet."""
+    if not types:
+        raise ValueError("types is required.")
+    return {
+        "trnm": "REMOVE",
+        "grp_no": group_no,
+        "data": [
+            {
+                "item": items,
+                "type": types,
+            }
+        ],
+    }