开云体育

# """ wrap connection.Connection to save TWS <-> ib_async messages, for later playback
# """

import atexit
import datetime
import inspect
import logging
import pickle
from functools import wraps
from typing import BinaryIO

import ib_async
from ib_async.connection import Connection

logger = logging.getLogger(__name__)


def _now():
    return datetime.datetime.now().astimezone()


# method decorator

def _pickle(method):

    _dump_name = method.__name__

    if inspect.iscoroutinefunction(method):

        @wraps(method)
        async def wrapper(self, *args):
            pickle.dump((_now(), _dump_name, *args), self._pickle_file)
            return await method(self, *args)

    else:

        @wraps(method)
        def wrapper(self, *args):
            pickle.dump((_now(), _dump_name, *args), self._pickle_file)
            return method(self, *args)

    return wrapper


class PickleConnection(Connection):
    _pickle_file: BinaryIO

    def __init__(self):
        self._open_pickle_file()
        super().__init__()

    def _open_pickle_file(self) -> None:
        timestamp = _now().strftime("%Y_%m_%d_%H_%M_%S")
        filename = f"connection_{timestamp}.pkl"
        atexit.register(self._close_pickle_file)
        self._pickle_file = open(filename, "wb+")

    def _close_pickle_file(self):
        if not self._pickle_file.closed:
            self._pickle_file.close()

    @_pickle
    async def connectAsync(self, host: str, port: int) -> None:
        await super().connectAsync(host, port)

    @_pickle
    def disconnect(self) -> None:
        super().disconnect()

    @_pickle
    def sendMsg(self, msg: bytes) -> None:
        super().sendMsg(msg)

    @_pickle
    def connection_lost(self, exc) -> None:
        super().connection_lost(exc)
        self._pickle_file.flush()

    @_pickle
    def data_received(self, data: bytes) -> None:
        super().data_received(data)

# patch ib_async to use PickleConnection instead of Connection

ib_async.connection.Connection = PickleConnection
ib_async.client.Connection = PickleConnection

# SCRIPT TO DUMP PKL FILE

# # !/usr/bin/env python3
#
# import datetime
# import sys
# import pickle
#
# filename = sys.argv[1]
#
# time: datetime.datetime
# method_name: str
#
# with open(filename, "rb") as fr:
#
#     try:
#         while True:
#             time, method_name, *args = pickle.load(fr)
#             timestamp = time.strftime("%H:%M:%S:%f")[:-1]
#             match method_name:
#                 case "data_received" | "sendMsg":
#                     assert len(args) == 1
#                     args = f"{len(args[0])}b"
#                 case _:
#                     args = ", ".join(map(str, args))
#             print(f"{timestamp}: {method_name}({args})")
#     except EOFError:
#         pass