Inconsistent typing for WebSocketMessage and WebSocketClient.run() callback in v2.8.0

まず、issue で言及されている WebSocketMessage、WebSocketClient.run()、parse_single()、parse() の定義を確認し、次にそれらの実行時の値と型シグネチャを調べます。WebSocketMessage が解析済みイベント 1 件を表すのか、バッチを表すのかを確認し、それに応じて callback と parser のアノテーションを整合させ、Pyright/Pylance が正しく型付けされた handler を受け入れられるように、run() に対して提案されているアノテーションを含めます。

Hi Massive team,

I noticed an inconsistency in the websocket typing in Massive Python SDK v2.8.0.

WebSocketMessage is currently defined as a NewType over a list of parsed event models:

WebSocketMessage = NewType(
    "WebSocketMessage",
    List[
        Union[
            EquityAgg,
            CurrencyAgg,
            EquityTrade,
            ...
        ]
    ],
)

However, WebSocketClient.run() is typed as accepting a callback of:

Callable[[List[WebSocketMessage]], None]

This effectively makes the callback parameter type a nested list-like structure:
List[WebSocketMessage], while WebSocketMessage itself already represents a list
of parsed websocket events.

This causes type checkers such as Pyright/Pylance to reject handlers like this:

def handle_msg(messages: WebSocketMessage) -> None:
    for message in messages:
        if isinstance(message, EquityAgg):
            ...

with an error similar to:

Argument of type "(messages: WebSocketMessage) -> None" cannot be assigned to parameter "handle_msg"
Type "(messages: WebSocketMessage) -> None" is not assignable to type "(List[WebSocketMessage]) -> None"

There also seems to be a related inconsistency in the parser:

def parse_single(...) -> Optional[WebSocketMessage]:
    parsed = model_class.from_dict(data)
    return cast(WebSocketMessage, parsed)

At runtime, parsed is a single event object such as EquityAgg, not a list. Then
parse() returns List[WebSocketMessage].

So there appear to be two possible fixes:

1. If WebSocketMessage is intended to mean a single parsed event, redefine it as
   a union of event model types, not a list.
2. If WebSocketMessage is intended to mean a batch of parsed events, then run()
   should accept Callable[[WebSocketMessage], None], and parse_single() should
   not cast individual event objects to WebSocketMessage.

Based on the current runtime behavior, option 1 seems more natural:

WebSocketMessage = NewType(
    "WebSocketMessage",
    Union[
        EquityAgg,
        CurrencyAgg,
        EquityTrade,
        ...
    ],
)

or alternatively, using a type alias instead of NewType:

WebSocketMessage = Union[
    EquityAgg,
    CurrencyAgg,
    EquityTrade,
    ...
]

Then these signatures would make sense:

def parse_single(...) -> Optional[WebSocketMessage]: ...
def parse(...) -> List[WebSocketMessage]: ...

def run(
    self,
    handle_msg: Callable[[List[WebSocketMessage]], None] | Callable[[str | bytes], None],
    close_timeout: int = 1,
    **kwargs: Any,
) -> None: ...

One smaller typing improvement: run() currently leaves **kwargs unknown for
Pyright/Pylance. Annotating it as **kwargs: Any and adding -> None would avoid
"partially unknown" warnings.

Thanks.