orderflow.py

import sys
import websockets
import asyncio
import json
import signal
from datetime import datetime
from datetime import timezone


class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    LIGHT_GREEN = '\033[42m'
    LIGHT_RED = '\033[101m'


class Trade:
    emoji = '💯'

    def __init__(self, market, time, side: str, size, price, symbol):
        self.market = market
        self.price = price
        self.size = size
        self.time = time
        self.side = side
        self.symbol = symbol

    def __str__(self):
        color = bcolors.LIGHT_RED
        if self.side.upper() == 'BUY':
            color = bcolors.LIGHT_GREEN
        px = self.price
        emoji_cnt = int(self.size / 100_000)

        return '[{symbol}] {date:%H:%M:%S} {color}{side:>4}{_end} {size:10,.1f} {price:7,.2f} {emoji}'.format(
            **{'color': color,
               'symbol': self.symbol,
               'date': self.time,
               'side': self.side.upper(),
               'size': self.size,
               'price': px,
               '_end': bcolors.ENDC,
               'emoji': (self.emoji * emoji_cnt)})

    def __repr__(self):
        return self.__str__()


stop = False


async def bitfinex(min_px=None):
    async with websockets.connect('wss://api.bitfinex.com/ws/2') as websocket:
        await websocket.send(json.dumps({
            'event': 'subscribe',
            'channel': 'trades',
            'symbol': 'BTCUSD'
        }))

        async for msg in websocket:
            if stop:
                break
            pkg = json.loads(msg)
            if 'event' in pkg or pkg[1] != 'te':
                continue

            channel, cmd, [seq, time, size, price] = pkg

            side = 'Buy'
            if size < 0:
                side = 'Sell'
            size = abs(size)
            if min_px and size < min_px:
                continue

            dt = datetime.fromtimestamp(time / 1000)
            t = Trade('BFX', dt, side, size, price, 'BTCUSD')
            yield t


async def bitmex(min_px=None):
    async with websockets.connect("wss://www.bitmex.com/realtime") as websocket:
        await websocket.send(json.dumps({
            'op': 'subscribe',
            'args': ['trade:XBTUSD', 'trade:XBTM18']
        }))
        async for msg in websocket:
            if stop:
                break
            pkg = json.loads(msg)
            if 'table' not in pkg or pkg['table'] != 'trade':
                continue
            for data_ in pkg['data']:
                if min_px and data_['size'] < min_px:
                    continue

                ts = datetime.strptime(data_['timestamp'], "%Y-%m-%dT%H:%M:%S.%fZ")
                ts = ts.replace(tzinfo=timezone.utc).astimezone()  # utc to local time
                t = Trade('MEX', ts, data_['side'], data_['size'], data_['price'], data_['symbol'])
                yield t


def sig_handler(a, b):
    global stop
    stop = True
    loop.stop()


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print('Provide min price parameter. Exiting.')
        sys.exit(1)


    def feed(min_px, retries=0):
        def wrapper():
            if len(sys.argv) == 3 and sys.argv[2] == 'bfx':
                return bitfinex
            return bitmex

        if retries >= 10:
            return []

        try:
            return wrapper()(min_px)
        except websockets.exceptions.ConnectionClosed as e:
            return feed(min_px, retries + 1)


    min_px = int(sys.argv[1])

    signal.signal(signal.SIGINT, sig_handler)
    loop = asyncio.get_event_loop()


    async def run():
        async for trade in feed(min_px):
            print(trade)


    # task = asyncio.ensure_future(run, loop=loop)
    loop.run_until_complete(run())