diff --git a/datamaxi/aio/ws.py b/datamaxi/aio/ws.py index c4c0199..fd3fb81 100644 --- a/datamaxi/aio/ws.py +++ b/datamaxi/aio/ws.py @@ -10,10 +10,14 @@ from datamaxi.aio.ws import AsyncDatamaxiWS async with AsyncDatamaxiWS(api_key="...") as ws: + # raw wire param ... async for msg in ws.ticker.subscribe("BTC-USDT@binance", market="spot"): print(msg["s"], msg["p"]) - async for oi in ws.open_interest.subscribe("BTC-USDT@binance"): + # ... or structured tokens named after the channel's param_format + async for oi in ws.open_interest.subscribe( + symbol="BTC-USDT", exchange="binance" + ): ... Channels (accessors driven by the generated ``WS_CHANNELS`` registry): @@ -217,6 +221,112 @@ def _require_channel(path: str) -> str: return path +def _token_kwarg(token: str) -> str: + """Map a raw format token to its kwarg name. + + Rule: lower-case a token only when it is entirely upper-case + (``SYMBOL`` -> ``symbol``); every other token is used verbatim + (``exchange``, ``tokenId``, ``srcQuote`` ...). + """ + return token.lower() if token.isupper() else token + + +def build_param(param_format: Optional[str], **tokens: str) -> str: + """Assemble a wire subscribe param from named tokens per a channel's format. + + ``param_format`` is the generated ``WS_CHANNELS[path]["param"]`` string, e.g. + ``"SYMBOL@exchange"`` or ``"src:tgt:tokenId:srcQuote:tgtQuote:srcMkt:tgtMkt"``. + Keyword names ARE the raw format tokens, with the single :func:`_token_kwarg` + rule (only an all-upper token is lower-cased, so ``SYMBOL`` -> ``symbol``; + ``exchange``/``tokenId``/``srcQuote`` stay verbatim). + + Grammar handled: + + * ``None`` -> the channel takes no params; structured kwargs are rejected. + * a single token (``SYMBOL``, forex) -> no separator. + * ``@``- or ``:``-separated tokens; the separator is inferred from the format. + * a trailing ``[...]`` group is optional and both-or-neither: supply every + token in it or none. + + Raises ``ValueError`` for an unknown token, a missing required token, or a + partially-supplied optional group; the message echoes ``param_format``. + + Examples:: + + build_param("SYMBOL", symbol="USD-KRW") # "USD-KRW" + build_param("SYMBOL@exchange", symbol="BTC-USDT", exchange="binance") + # -> "BTC-USDT@binance" + build_param( + "src:tgt:tokenId:srcQuote:tgtQuote:srcMkt:tgtMkt", + src="binance", tgt="upbit", tokenId="bitcoin", + srcQuote="USDT", tgtQuote="KRW", srcMkt="spot", tgtMkt="spot", + ) # -> "binance:upbit:bitcoin:USDT:KRW:spot:spot" + """ + if param_format is None: + raise ValueError( + "this channel takes no subscribe params; call subscribe() without " + "keyword tokens" + ) + + sep = "@" if "@" in param_format else ":" if ":" in param_format else "" + required_fmt, _, optional_fmt = param_format.partition("[") + optional_fmt = optional_fmt.rstrip("]") + + def _split(section: str) -> List[str]: + section = section.strip(sep) + if not section: + return [] + return section.split(sep) if sep else [section] + + required_kw = [_token_kwarg(t) for t in _split(required_fmt)] + optional_kw = [_token_kwarg(t) for t in _split(optional_fmt)] + valid = required_kw + optional_kw + + unknown = [k for k in tokens if k not in valid] + if unknown: + raise ValueError( + f"unknown subscribe token(s) {unknown} for format {param_format!r}; " + f"valid tokens: {valid}" + ) + + missing = [k for k in required_kw if k not in tokens] + if missing: + raise ValueError( + f"missing required subscribe token(s) {missing} for format " + f"{param_format!r}" + ) + + supplied_optional = [k for k in optional_kw if k in tokens] + if supplied_optional and len(supplied_optional) != len(optional_kw): + raise ValueError( + f"optional token group {optional_kw} is both-or-neither for format " + f"{param_format!r}; got only {supplied_optional}" + ) + + values = [str(tokens[k]) for k in required_kw + supplied_optional] + return sep.join(values) + + +def _resolve_params( + param_format: Optional[str], + params: tuple, + tokens: Dict[str, str], +) -> List[str]: + """Resolve a subscribe/unsubscribe call to a list of wire param strings. + + Raw positional ``params`` pass through unchanged (the backward-compatible + path). Keyword ``tokens`` build exactly one param via :func:`build_param`. + Mixing both is rejected. + """ + if params and tokens: + raise ValueError( + "pass either raw positional params or keyword tokens, not both" + ) + if tokens: + return [build_param(param_format, **tokens)] + return list(params) + + class Subscription: """A single-path subscribable channel (``ws.forex``, ``ws.premium``, ...). @@ -232,16 +342,27 @@ def __init__(self, client: "AsyncDatamaxiWS", path: str): def param_format(self) -> Optional[str]: return WS_CHANNELS[self._path].get("param") - async def subscribe(self, *params: str) -> AsyncIterator[Dict[str, Any]]: - """SUBSCRIBE to ``params``; return an async iterator over the channel.""" + async def subscribe( + self, *params: str, **tokens: str + ) -> AsyncIterator[Dict[str, Any]]: + """SUBSCRIBE and return an async iterator over the channel. + + Two forms: pass raw wire ``params`` positionally + (``subscribe("BTC-USDT@binance")``), or pass structured keyword + ``tokens`` named after this channel's :attr:`param_format` to build one + param (``subscribe(symbol="BTC-USDT", exchange="binance")``). See + :func:`build_param`. The two forms are mutually exclusive. + """ + resolved = _resolve_params(self.param_format, params, tokens) conn = await self._client._conn(self._path) stream = conn.stream() # register the queue before SUBSCRIBE (no missed msgs) - await conn.subscribe(list(params)) + await conn.subscribe(resolved) return stream - async def unsubscribe(self, *params: str) -> None: + async def unsubscribe(self, *params: str, **tokens: str) -> None: + resolved = _resolve_params(self.param_format, params, tokens) conn = await self._client._conn(self._path) - await conn.unsubscribe(list(params)) + await conn.unsubscribe(resolved) class MarketSubscription: @@ -259,16 +380,25 @@ def _path(self, market: str) -> str: return path async def subscribe( - self, *params: str, market: str = "spot" + self, *params: str, market: str = "spot", **tokens: str ) -> AsyncIterator[Dict[str, Any]]: - conn = await self._client._conn(self._path(market)) + """SUBSCRIBE on ``market``; ``params`` or structured ``tokens`` (see + :meth:`Subscription.subscribe`). ``market`` is a control kwarg, not a + param token.""" + path = self._path(market) + resolved = _resolve_params(WS_CHANNELS[path].get("param"), params, tokens) + conn = await self._client._conn(path) stream = conn.stream() - await conn.subscribe(list(params)) + await conn.subscribe(resolved) return stream - async def unsubscribe(self, *params: str, market: str = "spot") -> None: - conn = await self._client._conn(self._path(market)) - await conn.unsubscribe(list(params)) + async def unsubscribe( + self, *params: str, market: str = "spot", **tokens: str + ) -> None: + path = self._path(market) + resolved = _resolve_params(WS_CHANNELS[path].get("param"), params, tokens) + conn = await self._client._conn(path) + await conn.unsubscribe(resolved) class Feed: diff --git a/tests/test_ws.py b/tests/test_ws.py index da450a4..a0dea08 100644 --- a/tests/test_ws.py +++ b/tests/test_ws.py @@ -14,7 +14,7 @@ websockets = pytest.importorskip("websockets") import datamaxi._ws_models as _ws_models # noqa: E402 -from datamaxi.aio.ws import AsyncDatamaxiWS # noqa: E402 +from datamaxi.aio.ws import AsyncDatamaxiWS, build_param # noqa: E402 from datamaxi._ws_endpoints import WS_CHANNELS, WS_BASE_PATH # noqa: E402 from datamaxi._ws_models import TickerMessage # noqa: E402 @@ -281,3 +281,150 @@ async def run(): assert len(seen) >= 2 assert seen[0] == ["USD-KRW"] assert seen[1] == ["USD-KRW"] # `_open` replays sorted(self._active) + + +# --- structured subscribe helpers: build_param unit tests --- + + +def test_build_param_forex_single_token(): + # SYMBOL: no separator, one token; and the SYMBOL->symbol lowercase rule. + assert build_param("SYMBOL", symbol="USD-KRW") == "USD-KRW" + + +def test_build_param_symbol_exchange(): + assert ( + build_param("SYMBOL@exchange", symbol="BTC-USDT", exchange="binance") + == "BTC-USDT@binance" + ) + + +def test_build_param_ticker_without_optional_group(): + fmt = "SYMBOL@exchange[@currency@conversionBase]" + assert build_param(fmt, symbol="BTC-USDT", exchange="binance") == "BTC-USDT@binance" + + +def test_build_param_ticker_with_optional_group(): + fmt = "SYMBOL@exchange[@currency@conversionBase]" + assert ( + build_param( + fmt, + symbol="BTC-USDT", + exchange="binance", + currency="KRW", + conversionBase="USDT", + ) + == "BTC-USDT@binance@KRW@USDT" + ) + + +def test_build_param_premium_all_seven_tokens(): + fmt = "src:tgt:tokenId:srcQuote:tgtQuote:srcMkt:tgtMkt" + assert ( + build_param( + fmt, + src="binance", + tgt="upbit", + tokenId="bitcoin", + srcQuote="USDT", + tgtQuote="KRW", + srcMkt="spot", + tgtMkt="spot", + ) + == "binance:upbit:bitcoin:USDT:KRW:spot:spot" + ) + + +def test_build_param_verbatim_token_not_lowercased(): + # tokenId is mixed-case, so it stays verbatim (not lowercased). + with pytest.raises(ValueError): + build_param( + "src:tgt:tokenId:srcQuote:tgtQuote:srcMkt:tgtMkt", + src="a", + tgt="b", + tokenid="c", # wrong: lowercased key + srcQuote="d", + tgtQuote="e", + srcMkt="f", + tgtMkt="g", + ) + + +def test_build_param_none_channel_rejects_tokens(): + with pytest.raises(ValueError): + build_param(None, symbol="X") + + +def test_build_param_missing_required_token(): + with pytest.raises(ValueError) as ei: + build_param("SYMBOL@exchange", symbol="BTC-USDT") + assert "exchange" in str(ei.value) + + +def test_build_param_unknown_token(): + with pytest.raises(ValueError) as ei: + build_param("SYMBOL@exchange", symbol="BTC-USDT", exchange="binance", bogus="x") + assert "bogus" in str(ei.value) + + +def test_build_param_partial_optional_group_rejected(): + fmt = "SYMBOL@exchange[@currency@conversionBase]" + with pytest.raises(ValueError): + build_param(fmt, symbol="BTC-USDT", exchange="binance", currency="KRW") + + +def test_build_param_roundtrip_matches_raw_ticker(): + fmt = WS_CHANNELS["/ticker/spot"]["param"] + assert build_param(fmt, symbol="BTC-USDT", exchange="binance") == "BTC-USDT@binance" + + +def test_build_param_roundtrip_matches_raw_premium(): + fmt = WS_CHANNELS["/premium"]["param"] + assert ( + build_param( + fmt, + src="binance", + tgt="upbit", + tokenId="bitcoin", + srcQuote="USDT", + tgtQuote="KRW", + srcMkt="spot", + tgtMkt="spot", + ) + == "binance:upbit:bitcoin:USDT:KRW:spot:spot" + ) + + +def test_ws_structured_subscribe_produces_expected_wire_param(): + # A structured subscribe(...) call sends the same wire param as the raw form. + async def handler(conn): + async for raw in conn: + m = json.loads(raw) + if m.get("method") == "SUBSCRIBE": + await conn.send(json.dumps({"result": m["params"], "id": m["id"]})) + for p in m["params"]: + sym, exch = p.split("@")[0], p.split("@")[1] + await conn.send( + json.dumps({"s": sym, "e": exch, "p": 1.0, "d": 1, "_param": p}) + ) + + async def run(): + async with _serve(handler) as server: + async with AsyncDatamaxiWS( + api_key="k", ws_url=f"ws://localhost:{_port(server)}" + ) as ws: + stream = await ws.ticker.subscribe( + symbol="BTC-USDT", exchange="binance", market="spot" + ) + return await _first(stream) + + msg = _run(run()) + assert msg["_param"] == "BTC-USDT@binance" + + +def test_ws_structured_subscribe_mixing_raw_and_tokens_raises(): + async def run(): + async with AsyncDatamaxiWS(api_key="k", ws_url="ws://localhost:1") as ws: + await ws.forex.subscribe("USD-KRW", symbol="USD-KRW") + + with pytest.raises(ValueError): + _run(run())