mirror of
https://github.com/wassname/catalyst.git
synced 2026-06-28 13:00:23 +08:00
120 lines
3.0 KiB
Python
120 lines
3.0 KiB
Python
"""
|
|
Provides data handlers that can push messages to a zipline.core.DataFeed
|
|
|
|
::
|
|
DataSource
|
|
|
|
|
TradeDataSource
|
|
/ \
|
|
RandomEquityTrades SpecificEquityTrades
|
|
|
|
"""
|
|
import pytz
|
|
import random
|
|
import datetime
|
|
from mock import Mock
|
|
|
|
from zipline.components import DataSource
|
|
from zipline.utils import ndict
|
|
|
|
import zipline.protocol as zp
|
|
|
|
class TradeDataSource(DataSource):
|
|
|
|
def init(self):
|
|
self.setup_source()
|
|
|
|
#@property
|
|
#def get_id(self):
|
|
# return 'TradeDataSource'
|
|
|
|
def send(self, event):
|
|
"""
|
|
Sends the event iff it matches the internal sid filter.
|
|
:param dict event: is a trade event with data as per
|
|
:py:func: `zipline.protocol.TRADE_FRAME`
|
|
:rtype: None
|
|
"""
|
|
|
|
event.source_id = self.get_id
|
|
|
|
if event.sid in self.filter['sid']:
|
|
message = zp.DATASOURCE_FRAME(event)
|
|
self.data_socket.send(message)
|
|
|
|
|
|
class RandomEquityTrades(TradeDataSource):
|
|
"""
|
|
Generates a random stream of trades for testing.
|
|
"""
|
|
|
|
def init(self, sid, count):
|
|
self.count = count
|
|
self.incr = 0
|
|
self.sid = sid
|
|
self.trade_start = datetime.datetime.now().replace(tzinfo=pytz.utc)
|
|
self.day = datetime.timedelta(days=1)
|
|
self.price = random.uniform(5.0, 50.0)
|
|
|
|
self.setup_source()
|
|
|
|
@property
|
|
def get_id(self):
|
|
return 'RandomEquityTrades'
|
|
|
|
def do_work(self):
|
|
if not self.incr < self.count:
|
|
self.signal_done()
|
|
return
|
|
|
|
self.price = self.price + random.uniform(-0.05, 0.05)
|
|
volume = random.randrange(100,10000,100)
|
|
|
|
event = zp.ndict({
|
|
"type" : zp.DATASOURCE_TYPE.TRADE,
|
|
"sid" : self.sid,
|
|
"price" : self.price,
|
|
"volume" : volume,
|
|
"dt" : self.trade_start + (self.day * self.incr),
|
|
})
|
|
self.send(event)
|
|
self.incr += 1
|
|
|
|
|
|
class SpecificEquityTrades(TradeDataSource):
|
|
"""
|
|
Generates a non-random stream of trades for testing.
|
|
"""
|
|
|
|
def init(self, event_list):
|
|
"""
|
|
:param event_list: should be a chronologically ordered list of
|
|
dictionaries in the following form::
|
|
|
|
event = {
|
|
'sid' : an integer for security id,
|
|
'dt' : datetime object,
|
|
'price' : float for price,
|
|
'volume' : integer for volume
|
|
}
|
|
"""
|
|
self.event_list = event_list
|
|
self.count = 0
|
|
|
|
# TODO temporary hack
|
|
self.control_out = Mock()
|
|
self.setup_source()
|
|
|
|
@property
|
|
def get_id(self):
|
|
return "SpecificEquityTrades"
|
|
|
|
def do_work(self):
|
|
if(len(self.event_list) == 0):
|
|
self.signal_done()
|
|
return
|
|
|
|
event = self.event_list.pop(0)
|
|
self.send(zp.ndict(event))
|
|
self.count +=1
|