Plots candlestick chart from a JSON formatted tick data file

plotcandle.py

#!/home/alan/miniconda3/bin/python3

import sys
import json
from datetime import datetime
from recordclass import recordclass
from matplotlib import pyplot as plt


Candle = recordclass('Candle', ['open', 'close', 'hi', 'low', 'ts', 'volume'])


class CandleStick:
    '''Aggregator for tick data into CandleStick format'''
    def __init__(s, period):
        s.bars = []
        s.period = period
        s.last_ts = None

    def __getitem__(s, item):
        return s.bars[item]

    def __len__(s):
        return len(s.bars)

    def update(s, price, ts, volume=0, action=0):

        # Initialisation
        if s.last_ts is None:
            s.volume = volume
            s.last_ts = ts
            s.bars.append(Candle(price, price, price, price, int(ts/s.period), volume))
            return

        # Ignore ticks that arrived out of order
        if s.last_ts > ts:
            return

        # New candle
        if int(ts / s.period) != int(s.last_ts / s.period):
            ts_tmp = s.last_ts + s.period

            # append empty bars if tick arrived 1+ bar after the previous one
            while int(ts_tmp / s.period) < int(ts / s.period):

                # use close price of bar from last tick
                s.bars.append(Candle(
                        s.bars[-1].close,
                        s.bars[-1].close,
                        s.bars[-1].close,
                        s.bars[-1].close,
                        int(ts_tmp/s.period),
                        0)
                )

                ts_tmp = ts_tmp + s.period

            # append the new bar that contains the newly arrived tick
            s.bars.append(Candle(price, price, price, price, int(ts/s.period), volume))
            s.last_ts = ts

        # Old candle
        elif int(ts / s.period) == int(s.last_ts / s.period):
            s.bars[-1].low   = min(s.bars[-1].low, price)
            s.bars[-1].hi    = max(s.bars[-1].hi, price)
            s.bars[-1].close = price
            s.bars[-1].volume += volume


def plotCandles(candle, title=None):
    '''Plot candlestick charts based a list of candle tuples'''

    fig, ax1 = plt.subplots(1, 1, figsize=(16,12))

    if title:
        ax1.set_title(title)

    for i, bar in enumerate(candle):
        op = bar.open
        cl = bar.close
        hi = bar.hi
        lo = bar.low

        barlen = abs(op - cl)
        bottom = min(op, cl)
        color = 'r' if op > cl else 'g'
        fill  = True if color == 'r' else False

        ax1.bar(i, barlen, bottom=bottom, color=color, edgecolor=color, fill=fill, linewidth=1)
        ax1.vlines(i, lo, hi, color=color, linewidth=1)

    time_format = '%m-%d %H:%M'

    xlabels = [i * int(len(candle)/7.8) for i in range(8)]
    dt = [datetime.fromtimestamp(candle[x].ts * candle.period).strftime(time_format) for x in xlabels]
    plt.xticks(xlabels, dt)


if __name__ == '__main__':

    try:
        period = int(sys.argv[2])
    except:
        period = 900

    dataset = sys.argv[1]
    candle = CandleStick(period)

    with open(dataset) as fp:
        strings = fp.read().splitlines()
        ticks = [json.loads(s) for s in strings]
        for tick in ticks:
            candle.update(tick['price'], int(tick['timestamp']), tick['amount'])

    plotCandles(candle)
    plt.show()