bbgo_origin/python/bbgo/data/event.py

from __future__ import annotations

from dataclasses import dataclass
from typing import List

import bbgo_pb2

from ..enums import ChannelType
from ..enums import EventType
from .balance import Balance
from .depth import Depth
from .order import Order
from .trade import Trade


class Event:
    pass


@dataclass
class UserDataEvent(Event):
    session: str
    exchange: str
    channel_type: ChannelType
    event_type: EventType
    balances: List[Balance] = None
    trades: List[Trade] = None
    orders: List[Order] = None

    @classmethod
    def from_pb(cls, obj: bbgo_pb2.UserData) -> UserDataEvent:
        return cls(
            session=obj.session,
            exchange=obj.exchange,
            channel_type=ChannelType(obj.channel),
            event_type=EventType(obj.event),
            balances=[Balance.from_pb(balance) for balance in obj.balances],
            trades=[Trade.from_pb(trade) for trade in obj.trades],
            orders=[Order.from_pb(order) for order in obj.orders],
        )


@dataclass
class MarketDataEvent(Event):
    exchange: str
    symbol: str
    channel_type: ChannelType
    event_type: EventType
    depth: Depth = None

    @classmethod
    def from_pb(cls, obj: bbgo_pb2.MarketData) -> MarketDataEvent:
        channel_type = ChannelType(obj.channel)

        event = cls(
            exchange=obj.exchange,
            symbol=obj.symbol,
            channel_type=channel_type,
            event_type=EventType(obj.event),
        )

        if channel_type == ChannelType.BOOK:
            event.depth = Depth.from_pb(obj.depth)

        return event