Pythonノート11-VNに基づく.pyのリアルタイムKラインジェネレータ

from vnpy.trader.vtObject import VtBarData
from datetime import timedelta
import pymongo
from queue import Queue, Empty
from threading import Thread

class Bargenerate02:
    '''       K   '''
    def __init__(self, x=5):
        #   K       
        self.bar = None
        self.lastTick = None
        #   XminK 
        self.minxBar = None
        self.minX = x
        #        collection
        self.collections = {}
        #       
        self.active = False  #   
        self.queue = Queue()  #   
        self.thread = Thread(target=self.run)  #       
        #       
        self.callOpen = 0
        self.callAction = 0
        #          
        self.initVolume = 0

    def setDBname(self, dbname):
        '''        '''
        self.dbname = dbname

    def connetDB(self):
        '''     '''
        self.client = pymongo.MongoClient('localhost', 27017)

    def createCollection(self, symbols):
        '''
                    collection
          ：  K ：rb1905 X  K ：rb1905.Xmin
        :param symbols: list       
        :return:
        '''
        for symbol in symbols:
            self.collections[symbol] = []
            self.symbol = symbol
            self.symbolXmin = symbol + '.' + str(self.minX) + 'min'
            collection = self.client[self.dbname][self.symbol]
            collection.ensure_index([('datetime', pymongo.DESCENDING)])
            collectionXmin = self.client[self.dbname][self.symbolXmin]
            collectionXmin.ensure_index([('datetime', pymongo.DESCENDING)])
            self.collections[symbol].append(collection)
            self.collections[symbol].append(collectionXmin)

    def formatMin(self, dateTime):
        '''     '''
        #      ，10:55:50     10:56:00
        dateTime += timedelta(minutes=1)
        dateTime = dateTime.replace(second=0, microsecond=0)
        return dateTime

    def updateTick(self, tick):
        '''       tick   '''
        #   tick    
        self.barTime = tick.datetime
        if self.bar:
            if self.barTime.day != self.lastTick.datetime.day:
                #   ，     bar
                self.createNewBar(tick, first=1)

            elif self.barTime >= self.formatMin(self.lastTick.datetime):

                #                ，    tick     bar
                #        bar time，      bar
                self.updateBar(self.lastTick, timeFormat=1)
                if self.checkClose(self.barTime):
                    #         ，    tick，       
                    #       bar
                    self.updateBar(tick)
                else:
                    self.createNewBar(tick)

            else:
                #      ，  
                self.updateBar(tick)
        else:
            #   bar     
            #      bar
            self.createNewBar(tick, first=1)

    def checkCallAction(self, t):
        '''      '''
        if (t.hour == 12 and t.minute == 59) or (
                t.hour == 0 and t.minute == 59):
            return True

    def checkOpenLose(self, now, tick):
        '''            K '''
        hour = now.hour
        #      
        if 1 <= hour < 13:
            nightOpenTime = now.replace(hour=1, minute=0, second=0, microsecond=0)
            self.addBlankBar(now, nightOpenTime, tick, first=1)
        elif 13 <= hour < 19:
            dayOpenTime = now.replace(hour=13, minute=0, second=0, microsecond=0)
            self.addBlankBar(now, dayOpenTime, tick, first=1)

    def addBlankBar(self, now, last, tick, first=0):
        '''
            K 
                     ，     ，  K  OHLC        ，
                ，       。
                  ，  OHLC     tick      
        :param now:   tick  
        :param last:          
        :param tick:
        :param first:         
        :return:
        '''
        delta = (now - last).seconds // 60
        for i in range(1, delta + 1):
            bar_datetime = last + timedelta(minutes=i)
            if (bar_datetime.hour >= 3 and bar_datetime.hour < 13) \
                    or (bar_datetime.hour == 15 and bar_datetime.minute >= 30) \
                    or (bar_datetime.hour == 16) \
                    or (bar_datetime.hour == 17 and bar_datetime.minute <= 30) \
                    or (bar_datetime.hour >= 19) \
                    or (bar_datetime.hour == 14 and 15 <= bar_datetime.minute <= 30):
                continue
            bar = VtBarData()
            if self.callAction:
                #           ，       call_open
                bar.open = self.callOpen
                bar.high = self.callOpen
                bar.low = self.callOpen
                bar.close = self.callOpen
                #       
                #      ，       
                # self.callOpen = 0
                # self.callAction = 0
                #                 
                self.lastVolume = 0
            else:
                if first:
                    bar.open = tick.preClosePrice
                    bar.high = tick.preClosePrice
                    bar.low = tick.preClosePrice
                    bar.close = tick.preClosePrice
                else:
                    bar.open = tick.lastPrice
                    bar.high = tick.lastPrice
                    bar.low = tick.lastPrice
                    bar.close = tick.lastPrice
            bar.openInterest = tick.openInterest
            bar.volume = '0'
            bar.date = bar_datetime.strftime('%Y%m%d')
            bar.time = bar_datetime.strftime('%H:%M:%S')
            bar.datetime = bar_datetime
            bar.gatewayName = tick.gatewayName
            bar.symbol = tick.symbol
            bar.vtSymbol = tick.vtSymbol
            bar.exchange = tick.exchange
            self.insertData(bar, update=0)

    def checkClose(self, now):
        '''       '''
        huor = now.hour
        minute = now.minute
        second = now.second
        microsecond = now.microsecond
        if huor == 3 and minute == 0 and second == 0 and microsecond == 0:
            return True
        elif huor == 14 and minute == 15 and second == 0 and microsecond == 0:
            return True
        elif huor == 15 and minute == 30 and second == 0 and microsecond == 0:
            return True
        elif huor == 19 and minute == 0 and second == 0 and microsecond == 0:
            return True

    def insertData(self, data, index=0, update=1, timeFormat=0):
        '''
         tick  bar    
        :param data: tick，bar
        :param index: 0  1min 1   xmin
        :param update: 0  insert 1    update
        :return:
        '''
        dataPut = data
        self.queue.put((dataPut, index, update, timeFormat))

    def createNewBar(self, tick, first=0):
        '''    K '''
        #       ，        
        if self.checkCallAction(self.barTime):
            #     TICK  ，  lastPrice
            #     ，      TICK  ，     tick，        
            self.callOpen = tick.lastPrice
            self.callAction = 1
            self.lastTick = tick
        else:
            #        
            if first:
                #    ，      
                self.checkOpenLose(self.barTime, tick)
            else:
                self.addBlankBar(self.barTime, self.formatMin(self.lastTick.datetime), self.lastTick)
            #      bar
            self.bar = VtBarData()
            if self.callAction:
                #           ，       call_open
                self.bar.open = self.callOpen
                #       
                self.callOpen = 0
                self.callAction = 0
                #        ，        
                self.initVolume = 0
            else:
                self.bar.open = tick.lastPrice
                if self.lastTick:
                    if first:
                        #            
                        self.initVolume = 0
                    else:
                        #       tick，  volume     volume
                        self.initVolume = self.lastTick.volume
            self.bar.high = tick.lastPrice
            self.bar.low = tick.lastPrice
            self.bar.close = tick.lastPrice
            self.bar.openInterest = tick.openInterest
            self.bar.volume = int(tick.volume) - self.initVolume
            self.bar.date = tick.date
            self.bar.time = tick.time
            self.bar.datetime = tick.datetime
            self.bar.gatewayName = tick.gatewayName
            self.bar.symbol = tick.symbol
            self.bar.vtSymbol = tick.vtSymbol
            self.bar.exchange = tick.exchange

            self.insertData(self.bar, update=0)
            #      tick
            self.lastTick = tick

    def updateBar(self, tick, timeFormat=0):
        '''  K '''
        self.bar.high = max(self.bar.high, tick.lastPrice)
        self.bar.low = min(self.bar.low, tick.lastPrice)
        self.bar.close = tick.lastPrice
        self.bar.openInterest = tick.openInterest
        self.bar.volume = str(int(tick.volume) - self.initVolume)
        self.bar.time = tick.time
        self.bar.datetime = tick.datetime
        self.insertData(self.bar, timeFormat=timeFormat)
        self.lastTick = tick

    def run(self):
        '''       '''
        while self.active:
            try:
                data, index, update, timeFormat = self.queue.get(block=True, timeout=1)
                if update:
                    flt = {'_id': data._id}

                    if timeFormat:
                        data.datetime = self.formatMin(data.datetime)
                    self.collections[data.symbol][index].update(flt, {'$set': data.__dict__})
                else:
                    self.collections[data.symbol][index].insert(data.__dict__)
            except Empty:
                pass

    def start(self):
        '''  '''
        self.active = True
        self.thread.start()

    def stop(self):
        '''  '''
        if self.active:
            self.active = False
            self.thread.join()