Skip to content

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

Description

@aircozy

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions