mirror of
https://github.com/wassname/catalyst.git
synced 2026-09-12 12:12:04 +08:00
Reverted path independence code in factory.py. Not sure how to fix with new path structure.
This commit is contained in:
@@ -0,0 +1,5 @@
|
||||
from protocol_utils import ndict
|
||||
|
||||
__all__ = [
|
||||
ndict,
|
||||
]
|
||||
@@ -0,0 +1,130 @@
|
||||
from collections import namedtuple
|
||||
|
||||
import time
|
||||
import pytz
|
||||
import iso8601
|
||||
import calendar
|
||||
from dateutil import rrule
|
||||
from datetime import datetime, date, timedelta
|
||||
from dateutil.relativedelta import *
|
||||
|
||||
# Datetime Tuple
|
||||
# --------------
|
||||
d_tuple = namedtuple('dt', ['year', 'month', 'day', 'hour', 'minute', 'second', 'micros'])
|
||||
|
||||
# iso8061 utility
|
||||
# ---------------------
|
||||
def parse_iso8061(date_string):
|
||||
dt = iso8601.parse_date(date_string)
|
||||
dt = dt.replace(tzinfo = pytz.utc)
|
||||
return dt
|
||||
|
||||
# Epoch utilities
|
||||
# ---------------------
|
||||
UNIX_EPOCH = datetime(1970, 1, 1, 0, 0, tzinfo = pytz.utc)
|
||||
def EPOCH(utc_datetime):
|
||||
"""
|
||||
The key is to ensure all the dates you are using are in the utc timezone
|
||||
before you start converting. See http://pytz.sourceforge.net/ to learn how
|
||||
to do that properly. By normalizing to utc, you eliminate the ambiguity of
|
||||
daylight savings transitions. Then you can safely use timedelta to calculate
|
||||
distance from the unix epoch, and then convert to seconds or milliseconds.
|
||||
|
||||
Note that the resulting unix timestamp is itself in the UTC timezone. If you
|
||||
wish to see the timestamp in a localized timezone, you will need to make
|
||||
another conversion.
|
||||
|
||||
Also note that this will only work for dates after 1970.
|
||||
"""
|
||||
assert isinstance(utc_datetime, datetime)
|
||||
# utc only please
|
||||
assert utc_datetime.tzinfo == pytz.utc
|
||||
|
||||
# how long since the epoch?
|
||||
delta = utc_datetime - UNIX_EPOCH
|
||||
seconds = delta.total_seconds()
|
||||
ms = seconds * 1000
|
||||
return ms
|
||||
|
||||
def UN_EPOCH(ms_since_epoch):
|
||||
seconds_since_epoch = ms_since_epoch / 1000
|
||||
delta = timedelta(seconds = seconds_since_epoch)
|
||||
dt = UNIX_EPOCH + delta
|
||||
return dt
|
||||
|
||||
def iso8061_to_epoch(datestring):
|
||||
dt = parse_iso8061(datestring)
|
||||
return EPOCH(dt)
|
||||
|
||||
def epoch_now():
|
||||
dt = datetime.utcnow().replace(tzinfo=pytz.utc)
|
||||
return EPOCH(dt)
|
||||
|
||||
# UTC Datetime Subclasses
|
||||
# -----------------------
|
||||
def utcnow():
|
||||
return datetime.now(pytz.utc)
|
||||
|
||||
class utcdatetime(datetime):
|
||||
def __new__(cls, *args, **kwargs):
|
||||
kwargs['tzinfo'] = pytz.utc
|
||||
dt = datetime.__new__(cls, *args, **kwargs)
|
||||
return dt
|
||||
|
||||
|
||||
|
||||
# Datetime Calculations
|
||||
# ---------------------
|
||||
|
||||
WEEKDAYS = [rrule.MO, rrule.TU, rrule.WE, rrule.TH, rrule.FR]
|
||||
|
||||
HOLIDAYS = {
|
||||
'new_years' : datetime(2008 , 1 , 1 ),
|
||||
'mlk_day' : datetime(2008 , 1 , 21),
|
||||
'presidents' : datetime(2008 , 2 , 18),
|
||||
'good_friday' : datetime(2008 , 3 , 21),
|
||||
'memorial_day' : datetime(2008 , 5 , 26),
|
||||
'july_4th' : datetime(2008 , 7 , 4 ),
|
||||
'labor_day' : datetime(2008 , 9 , 1 ),
|
||||
'tgiving' : datetime(2008 , 11 , 27),
|
||||
'christmas' : datetime(2008 , 5 , 25),
|
||||
}
|
||||
|
||||
# Create a rule to recur every weekday starting today
|
||||
rule = rrule.rrule(
|
||||
rrule.DAILY,
|
||||
byweekday=WEEKDAYS,
|
||||
cache=True,
|
||||
)
|
||||
|
||||
# Precompute the rule, so that dates are cached.
|
||||
rs = rrule.rruleset()
|
||||
rs.rrule(rule)
|
||||
|
||||
# Add holidays as exclusion days
|
||||
for holiday in HOLIDAYS.itervalues():
|
||||
rs.exdate(holiday)
|
||||
|
||||
def trading_days(after, before, inclusive=False):
|
||||
"""
|
||||
Iterates over the NYSE trading days between the two given
|
||||
dates.
|
||||
"""
|
||||
return rs.between(after, before, inc=inclusive)
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
now = datetime.now()
|
||||
now30 = datetime.now() + timedelta(days=30)
|
||||
|
||||
# Iterate over the trading days between any two arbitrary
|
||||
# days, excluding the preset holidays.
|
||||
for day in trading_days(now, now30):
|
||||
print day
|
||||
|
||||
# Its now cached so if we do that traversal again it only
|
||||
# takes like 1e-5 seconds.
|
||||
tic = time.time()
|
||||
for day in trading_days(now, now30):
|
||||
print day
|
||||
print time.time() - tic
|
||||
@@ -0,0 +1,236 @@
|
||||
"""
|
||||
Factory functions to prepare useful data for tests.
|
||||
"""
|
||||
|
||||
import pytz
|
||||
import msgpack
|
||||
import random
|
||||
from os.path import join
|
||||
from operator import attrgetter
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
import zipline.finance.risk as risk
|
||||
import zipline.protocol as zp
|
||||
from zipline.finance.sources import SpecificEquityTrades, RandomEquityTrades
|
||||
from zipline.finance.trading import TradingEnvironment
|
||||
|
||||
def load_market_data():
|
||||
fp_bm = open("./tests/benchmark.msgpack", "rb")
|
||||
bm_list = msgpack.loads(fp_bm.read())
|
||||
bm_returns = []
|
||||
for packed_date, returns in bm_list:
|
||||
event_dt = zp.tuple_to_date(packed_date)
|
||||
#event_dt = event_dt.replace(
|
||||
# hour=0,
|
||||
# minute=0,
|
||||
# second=0,
|
||||
# tzinfo=pytz.utc
|
||||
#)
|
||||
|
||||
daily_return = risk.DailyReturn(date=event_dt, returns=returns)
|
||||
bm_returns.append(daily_return)
|
||||
|
||||
bm_returns = sorted(bm_returns, key=attrgetter('date'))
|
||||
fp_tr = open(".//tests/treasury_curves.msgpack", "rb")
|
||||
tr_list = msgpack.loads(fp_tr.read())
|
||||
tr_curves = {}
|
||||
for packed_date, curve in tr_list:
|
||||
tr_dt = zp.tuple_to_date(packed_date)
|
||||
#tr_dt = tr_dt.replace(hour=0, minute=0, second=0, tzinfo=pytz.utc)
|
||||
tr_curves[tr_dt] = curve
|
||||
|
||||
return bm_returns, tr_curves
|
||||
|
||||
def create_trading_environment(year=2006):
|
||||
"""Construct a complete environment with reasonable defaults"""
|
||||
benchmark_returns, treasury_curves = load_market_data()
|
||||
|
||||
start = datetime(year, 1, 1, tzinfo=pytz.utc)
|
||||
end = datetime(year, 12, 31, tzinfo=pytz.utc)
|
||||
trading_environment = TradingEnvironment(
|
||||
benchmark_returns,
|
||||
treasury_curves,
|
||||
period_start = start,
|
||||
period_end = end,
|
||||
capital_base = 100000.0
|
||||
)
|
||||
|
||||
return trading_environment
|
||||
|
||||
def create_trade(sid, price, amount, datetime):
|
||||
row = zp.ndict({
|
||||
'source_id' : "test_factory",
|
||||
'type' : zp.DATASOURCE_TYPE.TRADE,
|
||||
'sid' : sid,
|
||||
'dt' : datetime,
|
||||
'price' : price,
|
||||
'volume' : amount
|
||||
})
|
||||
return row
|
||||
|
||||
def get_next_trading_dt(current, interval, trading_calendar):
|
||||
next = current
|
||||
while True:
|
||||
next = next + interval
|
||||
if trading_calendar.is_market_hours(next):
|
||||
break
|
||||
|
||||
return next
|
||||
|
||||
def create_trade_history(sid, prices, amounts, interval, trading_calendar):
|
||||
trades = []
|
||||
current = trading_calendar.first_open
|
||||
|
||||
for price, amount in zip(prices, amounts):
|
||||
|
||||
trade = create_trade(sid, price, amount, current)
|
||||
trades.append(trade)
|
||||
current = get_next_trading_dt(current, interval, trading_calendar)
|
||||
|
||||
assert len(trades) == len(prices)
|
||||
return trades
|
||||
|
||||
def create_txn(sid, price, amount, datetime, btrid=None):
|
||||
txn = zp.ndict({
|
||||
'sid' : sid,
|
||||
'amount' : amount,
|
||||
'dt' : datetime,
|
||||
'price' : price,
|
||||
})
|
||||
return txn
|
||||
|
||||
def create_txn_history(sid, priceList, amtList, interval, trading_calendar):
|
||||
txns = []
|
||||
current = trading_calendar.first_open
|
||||
|
||||
for price, amount in zip(priceList, amtList):
|
||||
current = get_next_trading_dt(current, interval, trading_calendar)
|
||||
|
||||
txns.append(create_txn(sid, price, amount, current))
|
||||
current = current + interval
|
||||
return txns
|
||||
|
||||
|
||||
def create_returns(daycount, trading_calendar):
|
||||
"""
|
||||
For the given number of calendar (not trading) days return all the trading
|
||||
days between start and start + daycount.
|
||||
"""
|
||||
test_range = []
|
||||
current = trading_calendar.first_open
|
||||
one_day = timedelta(days = 1)
|
||||
|
||||
for day in range(daycount):
|
||||
current = current + one_day
|
||||
if trading_calendar.is_trading_day(current):
|
||||
r = risk.DailyReturn(current, random.random())
|
||||
test_range.append(r)
|
||||
|
||||
return test_range
|
||||
|
||||
|
||||
def create_returns_from_range(trading_calendar):
|
||||
current = trading_calendar.first_open
|
||||
end = trading_calendar.last_close
|
||||
one_day = timedelta(days = 1)
|
||||
test_range = []
|
||||
while current <= end:
|
||||
r = risk.DailyReturn(current, random.random())
|
||||
test_range.append(r)
|
||||
current = get_next_trading_dt(current, one_day, trading_calendar)
|
||||
|
||||
return test_range
|
||||
|
||||
def create_returns_from_list(returns, trading_calendar):
|
||||
current = trading_calendar.first_open
|
||||
one_day = timedelta(days = 1)
|
||||
test_range = []
|
||||
|
||||
#sometimes the range starts with a non-trading day.
|
||||
if not trading_calendar.is_trading_day(current):
|
||||
current = get_next_trading_dt(current, one_day, trading_calendar)
|
||||
|
||||
for return_val in returns:
|
||||
r = risk.DailyReturn(current, return_val)
|
||||
test_range.append(r)
|
||||
current = get_next_trading_dt(current, one_day, trading_calendar)
|
||||
|
||||
return test_range
|
||||
|
||||
def create_random_trade_source(sid, trade_count, trading_environment):
|
||||
# create the source
|
||||
source = RandomEquityTrades(sid, "rand-"+str(sid), trade_count)
|
||||
|
||||
# make the period_end of trading_environment match
|
||||
cur = trading_environment.first_open
|
||||
one_day = timedelta(days = 1)
|
||||
for i in range(trade_count + 2):
|
||||
cur = get_next_trading_dt(cur, one_day, trading_environment)
|
||||
trading_environment.period_end = cur
|
||||
|
||||
return source
|
||||
|
||||
def create_daily_trade_source(sids, trade_count, trading_environment):
|
||||
|
||||
"""
|
||||
creates trade_count trades for each sid in sids list.
|
||||
first trade will be on trading_environment.period_start, and daily
|
||||
thereafter for each sid. Thus, two sids should result in two trades per
|
||||
day.
|
||||
|
||||
Important side-effect: trading_environment.period_end will be modified
|
||||
to match the day of the final trade.
|
||||
"""
|
||||
return create_trade_source(
|
||||
sids,
|
||||
trade_count,
|
||||
timedelta(days=1),
|
||||
trading_environment
|
||||
)
|
||||
|
||||
|
||||
def create_minutely_trade_source(sids, trade_count, trading_environment):
|
||||
|
||||
"""
|
||||
creates trade_count trades for each sid in sids list.
|
||||
first trade will be on trading_environment.period_start, and every minute
|
||||
thereafter for each sid. Thus, two sids should result in two trades per
|
||||
minute.
|
||||
|
||||
Important side-effect: trading_environment.period_end will be modified
|
||||
to match the day of the final trade.
|
||||
"""
|
||||
return create_trade_source(
|
||||
sids,
|
||||
trade_count,
|
||||
timedelta(minutes=1),
|
||||
trading_environment
|
||||
)
|
||||
|
||||
def create_trade_source(sids, trade_count, trade_time_increment, trading_environment):
|
||||
trade_history = []
|
||||
|
||||
price = [10.1] * trade_count
|
||||
volume = [100] * trade_count
|
||||
|
||||
for sid in sids:
|
||||
start_date = trading_environment.first_open
|
||||
|
||||
generated_trades = create_trade_history(
|
||||
sid,
|
||||
price,
|
||||
volume,
|
||||
trade_time_increment,
|
||||
trading_environment
|
||||
)
|
||||
|
||||
trade_history.extend(generated_trades)
|
||||
|
||||
trade_history = sorted(trade_history, key=attrgetter('dt'))
|
||||
|
||||
#set the trading environment's end to same dt as the last trade in the
|
||||
#history.
|
||||
trading_environment.period_end = trade_history[-1].dt
|
||||
|
||||
source = SpecificEquityTrades("flat", trade_history)
|
||||
return source
|
||||
@@ -0,0 +1,99 @@
|
||||
"""
|
||||
This is somewhat legally ambigious, since it technically
|
||||
hasn't been merged in gevent_zeromq but given that the
|
||||
author issued it as a Pull Request on a MIT project,
|
||||
indicates that its probably fine to use. ~Steve
|
||||
"""
|
||||
|
||||
import zmq
|
||||
from zmq import *
|
||||
|
||||
from zmq.core.poll import Poller as _original_Poller
|
||||
|
||||
import gevent
|
||||
from gevent import select
|
||||
from gevent_zeromq.core import _Socket
|
||||
|
||||
def patch_poller(self):
|
||||
zmq.Poller = _Poller
|
||||
|
||||
class _Poller(_original_Poller):
|
||||
"""
|
||||
Replacement for :class:`zmq.core.Poller`
|
||||
|
||||
Ensures that the greened Poller below is used in calls
|
||||
to :meth:`zmq.core.Poller.poll`.
|
||||
"""
|
||||
|
||||
def _get_descriptors(self):
|
||||
"""
|
||||
Returns three elements tuple with socket descriptors ready for
|
||||
gevent.select
|
||||
"""
|
||||
rlist = []
|
||||
wlist = []
|
||||
xlist = []
|
||||
|
||||
for socket, flags in self.sockets.items():
|
||||
if isinstance(socket, _Socket):
|
||||
fd = socket.getsockopt(FD)
|
||||
elif isinstance(socket, int):
|
||||
fd = socket
|
||||
elif hasattr(socket, 'fileno'):
|
||||
try:
|
||||
fd = int(socket.fileno())
|
||||
except:
|
||||
raise ValueError('fileno() must return an valid integer fd')
|
||||
else:
|
||||
raise TypeError("Socket must be a 0MQ socket, an integer fd or \
|
||||
have a fileno() method: %r" % socket)
|
||||
|
||||
if flags & POLLIN: rlist.append(fd)
|
||||
if flags & POLLOUT: wlist.append(fd)
|
||||
if flags & POLLERR: xlist.append(fd)
|
||||
|
||||
return (rlist, wlist, xlist)
|
||||
|
||||
def poll(self, timeout=-1):
|
||||
"""Overridden method to ensure that the green version of Poller is used
|
||||
|
||||
Behaves the same as :meth:`zmq.core.Poller.poll`
|
||||
"""
|
||||
|
||||
if timeout is None:
|
||||
timeout = -1
|
||||
|
||||
timeout = int(timeout)
|
||||
if timeout < 0:
|
||||
timeout = -1
|
||||
|
||||
rlist = None
|
||||
wlist = None
|
||||
xlist = None
|
||||
|
||||
if timeout > 0:
|
||||
tout = gevent.Timeout.start_new(timeout/1000.0)
|
||||
|
||||
try:
|
||||
# Loop until timeout or events available
|
||||
while True:
|
||||
events = super(_Poller, self).poll(0)
|
||||
if events or timeout == 0:
|
||||
return events
|
||||
|
||||
# wait for activity on sockets in a green way
|
||||
if not rlist and not wlist and not xlist:
|
||||
rlist, wlist, xlist = self._get_descriptors()
|
||||
|
||||
try:
|
||||
select.select(rlist, wlist, xlist)
|
||||
except gevent.select.error, ex:
|
||||
raise ZMQError(*ex.args)
|
||||
|
||||
except gevent.Timeout, t:
|
||||
if t is not tout:
|
||||
raise
|
||||
return []
|
||||
finally:
|
||||
if timeout > 0:
|
||||
tout.cancel()
|
||||
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
Small classes to assist with timezone calculations, LOGGER configuration,
|
||||
and other common operations.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import logging.config
|
||||
|
||||
def configure_logging():
|
||||
logging.config.fileConfig(
|
||||
'logging.cfg',
|
||||
disable_existing_loggers = False
|
||||
)
|
||||
@@ -0,0 +1,174 @@
|
||||
import copy
|
||||
import pandas
|
||||
from ctypes import Structure, c_ubyte
|
||||
from collections import MutableMapping
|
||||
from itertools import izip
|
||||
|
||||
def Enum(*options):
|
||||
"""
|
||||
Fast enums are very important when we want really tight zmq
|
||||
loops. These are probably going to evolve into pure C structs
|
||||
anyways so might as well get going on that.
|
||||
"""
|
||||
class cstruct(Structure):
|
||||
_fields_ = [(o, c_ubyte) for o in options]
|
||||
__iter__ = lambda s: iter(range(len(options)))
|
||||
return cstruct(*range(len(options)))
|
||||
|
||||
def FrameExceptionFactory(name):
|
||||
"""
|
||||
Exception factory with a closure around the frame class name.
|
||||
"""
|
||||
class InvalidFrame(Exception):
|
||||
def __init__(self, got):
|
||||
self.got = got
|
||||
|
||||
def __str__(self):
|
||||
return "Invalid {framecls} Frame: {got}".format(
|
||||
framecls = name,
|
||||
got = self.got,
|
||||
)
|
||||
|
||||
return InvalidFrame
|
||||
|
||||
class ndict(MutableMapping):
|
||||
"""
|
||||
Xtreme Namedicts 2.0
|
||||
|
||||
Ndicts are dict like objects that have fields accessible by attribute
|
||||
lookup as well as being indexable and iterable. Done right
|
||||
this time.
|
||||
"""
|
||||
|
||||
def __init__(self, dct=None):
|
||||
self.__internal = dict()
|
||||
self.cls = frozenset(dir(self))
|
||||
|
||||
if dct:
|
||||
self.__internal.update(dct)
|
||||
|
||||
# Abstact Overloads
|
||||
# -----------------
|
||||
|
||||
def __setattr__(self, key, value):
|
||||
if '_ndict' in key or key == 'cls':
|
||||
self.__dict__[key] = value
|
||||
else:
|
||||
self.__internal[key] = value
|
||||
return value
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
"""
|
||||
Required for use by pymongo as_class parameter to find.
|
||||
"""
|
||||
if key == '_id':
|
||||
self.__internal['id'] = value
|
||||
else:
|
||||
self.__internal[key] = value
|
||||
|
||||
def __getattr__(self, key):
|
||||
if key in self.cls:
|
||||
return self.__dict__[key]
|
||||
else:
|
||||
return self.__internal[key]
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self.__internal[key]
|
||||
|
||||
def __delitem__(self, key):
|
||||
del self.__internal[key]
|
||||
|
||||
def __iter__(self):
|
||||
return self.__internal.iterkeys()
|
||||
|
||||
def __len__(self):
|
||||
return len(self.__internal)
|
||||
|
||||
# Compatability with namedicts
|
||||
# ----------------------------
|
||||
|
||||
# for compat, not the Python way to do things though...
|
||||
# Deprecated, use builtin ``del`` operator.
|
||||
delete = __delitem__
|
||||
|
||||
def has_attr(self, key):
|
||||
"""
|
||||
Deprecated, use builtin ``in`` operator.
|
||||
"""
|
||||
return self.__contains__(key)
|
||||
|
||||
def has_key(self, key):
|
||||
return self.__contains__(key)
|
||||
|
||||
# Custom Methods
|
||||
# --------------
|
||||
|
||||
def copy(self):
|
||||
return ndict(copy.copy(self.__internal))
|
||||
|
||||
def as_dataframe(self):
|
||||
"""
|
||||
Return the representation as a Pandas dataframe.
|
||||
"""
|
||||
d = pandas.DataFrame(self.__internal)
|
||||
return d
|
||||
|
||||
def as_series(self):
|
||||
"""
|
||||
Return the representation as a Pandas time series.
|
||||
"""
|
||||
s = pandas.Series(self.__internal)
|
||||
s.name = self.sid
|
||||
return s
|
||||
|
||||
def as_dict(self):
|
||||
"""
|
||||
Return the representation as a vanilla Python dict.
|
||||
"""
|
||||
# shallow copy is O(n)
|
||||
return copy.copy(self.__internal)
|
||||
|
||||
def merge(self, other_nd):
|
||||
"""
|
||||
Merge in place with another ndict.
|
||||
"""
|
||||
assert isinstance(other_nd, ndict)
|
||||
self.__internal.update(other_nd.__internal)
|
||||
|
||||
def __repr__(self):
|
||||
return "namedict: " + str(self.__internal)
|
||||
|
||||
# Faster dictionary comparison?
|
||||
#def __eq__(self, other):
|
||||
#assert isinstance(other, ndict)
|
||||
|
||||
#keyeq = set(self.keys()) == set(other.keys())
|
||||
|
||||
#if not keyeq:
|
||||
#return False
|
||||
|
||||
#for i, j in izip(self.itervalues(), other.itervalues()):
|
||||
#if i != j:
|
||||
#return False
|
||||
|
||||
#return True
|
||||
|
||||
# This is not neccesarily the most intuitive construction, but
|
||||
# we're aiming for raw performance rather than readability. So
|
||||
# we do things that we would not normally do in business logic.
|
||||
def namelookup(dct):
|
||||
ks = dct.keys()
|
||||
vs = dct.values()
|
||||
dct = {}
|
||||
class _lookup:
|
||||
__slots__ = ks
|
||||
def __init__(self):
|
||||
for k, v in zip(ks, vs):
|
||||
setattr(self,k,v)
|
||||
self.__setattr__ = self.locked
|
||||
def locked(self,k,v):
|
||||
raise Exception('Name lookups are fixed at init.')
|
||||
def __repr__(self):
|
||||
return '<namelookup %s>' % self.__slots__
|
||||
del dct
|
||||
return _lookup()
|
||||
@@ -0,0 +1,56 @@
|
||||
"""
|
||||
Format serializer for Zipline.
|
||||
|
||||
Because I'm opinionated about how you should send things over
|
||||
ZeroMQ. :)
|
||||
"""
|
||||
|
||||
import zlib
|
||||
import hmac
|
||||
import base64
|
||||
#import blosc
|
||||
|
||||
import cPickle as pickle
|
||||
|
||||
# Pickle does the equivelant of builtin ``eval``. Be afraid, be
|
||||
# very afraid.
|
||||
|
||||
def send_zipped_pickle(socket, obj, flags=0, protocol=-1):
|
||||
"""
|
||||
Pickle an object, and zip the pickle before sending it.
|
||||
"""
|
||||
p = pickle.dumps(obj, protocol)
|
||||
z = zlib.compress(p)
|
||||
return socket.send(z, flags=flags)
|
||||
|
||||
def recv_zipped_pickle(socket, flags=0, protocol=-1):
|
||||
"""
|
||||
Unpickle and uncompress a received object.
|
||||
"""
|
||||
z = socket.recv(flags)
|
||||
p = zlib.uncompress(z)
|
||||
return pickle.loads(p, protocol=protocol)
|
||||
|
||||
# Cryptographically secure wire protocol for ZeroMQ Using HMAC.
|
||||
|
||||
# Compare byte strings, backported from Python 3.
|
||||
def byte_eq(a, b):
|
||||
return not sum(0 if x==y else 1 for x, y in zip(a, b)) and len(a) == len(b)
|
||||
|
||||
def send_secure(socket, data, key, flags=0):
|
||||
msg = base64.b64encode(data)
|
||||
sig = base64.b64encode(hmac.new(key, msg).digest())
|
||||
return socket.send(bytes('!') + sig + bytes('?') + msg, flags=flags)
|
||||
|
||||
def recv_secure(socket, data, key, flags):
|
||||
data = socket.recv(flags=flags)
|
||||
|
||||
try:
|
||||
sig, msg = data.split(bytes('?'), 1)
|
||||
except ValueError:
|
||||
raise Exception('Invalid signature/message pair.')
|
||||
|
||||
if byte_eq(sig[1:], base64.b64encode(hmac.new(key, msg).digest())):
|
||||
return base64.b64decode(msg)
|
||||
else:
|
||||
raise Exception('Cryptographically invalid message received')
|
||||
@@ -0,0 +1,154 @@
|
||||
"""
|
||||
Misc ZeroMQ utilities.
|
||||
"""
|
||||
import gevent
|
||||
import msgpack
|
||||
import numpy
|
||||
from numpy import dtype
|
||||
from pandas import DataFrame
|
||||
from gevent_zeromq import zmq
|
||||
|
||||
from contextlib import closing
|
||||
|
||||
class ZmqDone(Exception):
|
||||
|
||||
def __init__(self, socket, frame):
|
||||
self.ident = socket.identity
|
||||
self.frame = str(frame)
|
||||
|
||||
def __str__(self):
|
||||
return 'Socket ( %s ) finished with frame ( %s )' % \
|
||||
( self.ident, self.frame )
|
||||
|
||||
class zs(object):
|
||||
"""
|
||||
A wrapper for the *very* common pattern of reading from a
|
||||
upstream socket until you get a DONE or EXCEPTION frame.
|
||||
|
||||
# Eliminates all the boilerplate serialization logic
|
||||
# and error handling cases into 3 lines.
|
||||
|
||||
halts = (ERROR_FRAME, CLOSE_FRAME)
|
||||
stream = zs(socket, halts)
|
||||
|
||||
stream.on_error(YouFailAtFailing)
|
||||
|
||||
for msg in stream:
|
||||
print msg
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, socket, halts, srl=msgpack):
|
||||
self._socket = socket
|
||||
self.exc_case = halts[0]
|
||||
self.done_case = halts[1]
|
||||
|
||||
self.loads = srl.loads
|
||||
self.halt_method = 'exception'
|
||||
self.exception = ZmqDone
|
||||
self.function = None
|
||||
|
||||
def __iter__(self):
|
||||
self.last = msg = self.loads(self._socket.recv())
|
||||
|
||||
if msg == self.exc_case:
|
||||
return self.halt()
|
||||
|
||||
if msg == self.done_case:
|
||||
raise StopIteration
|
||||
|
||||
yield msg
|
||||
|
||||
def last(self):
|
||||
return self.last
|
||||
|
||||
def halt(self):
|
||||
if self.halt_method == 'exception':
|
||||
raise self.exception
|
||||
elif self.halt_method == 'function':
|
||||
return self.function()
|
||||
|
||||
def on_error(self, callee):
|
||||
|
||||
if isinstance(callee, Exception):
|
||||
self.halt_method = 'exception'
|
||||
self.exception = callee
|
||||
else:
|
||||
self.halt_method = 'function'
|
||||
self.function = callee
|
||||
|
||||
def ZmqConsole(sock_typ, socket_addr, sock_conn=None, context=None):
|
||||
"""
|
||||
A utility to drop into a ZeroMQ pdb console and inspect
|
||||
messages as they come through. If you just want to pipe to
|
||||
stdout, don't use this.
|
||||
"""
|
||||
|
||||
context = context or zmq.Context.instance()
|
||||
socket = context.socket(zmq.PULL)
|
||||
socket.bind(socket_addr)
|
||||
|
||||
def console():
|
||||
while True:
|
||||
msg = socket.recv_pyobj()
|
||||
print msg
|
||||
import pdb; pdb.set_trace()
|
||||
|
||||
return gevent.spawn(console)
|
||||
|
||||
class NumpyChannel(zmq.Socket):
|
||||
|
||||
def recv_pandas(self, flags=0, copy=True, track=False):
|
||||
|
||||
# Pandas Metadata
|
||||
index, columns, dtype_name, shape = msgpack.loads(self.recv(flags=flags))
|
||||
|
||||
# Pandas ndarray
|
||||
ndbuffer = self.recv(flags=flags, copy=copy, track=track)
|
||||
buf = buffer(ndbuffer)
|
||||
|
||||
ndarray = numpy.frombuffer(buf, dtype=dtype(dtype_name)).reshape(shape)
|
||||
return DataFrame(data=ndarray, index=index,
|
||||
columns=columns, dtype=dtype_name)
|
||||
|
||||
def send_pandas(self, df, flags=0, copy=True, track=False):
|
||||
|
||||
# Pandas Metadata
|
||||
index = df.index.tolist()
|
||||
columns = df.columns.tolist()
|
||||
dtype_name = df.values.dtype.name
|
||||
shape = df.values.shape
|
||||
|
||||
# Pandas ndarray
|
||||
ndarray = df.values
|
||||
|
||||
metadata = msgpack.dumps((index, columns, dtype_name, shape))
|
||||
|
||||
self.send(metadata, flags|zmq.SNDMORE)
|
||||
return self.send(ndarray, flags, copy=copy, track=track)
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
from numpy.random import randn
|
||||
df = DataFrame(randn(5,5))
|
||||
|
||||
ctx = zmq.Context.instance()
|
||||
|
||||
def send():
|
||||
pub = NumpyChannel(ctx, zmq.PUSH)
|
||||
pub.bind('inproc://a')
|
||||
|
||||
for i in xrange(100):
|
||||
pub.send_pandas(df, copy=False)
|
||||
|
||||
def recv():
|
||||
sub = NumpyChannel(ctx, zmq.PULL)
|
||||
sub.connect('inproc://a')
|
||||
|
||||
for i in xrange(100):
|
||||
sub.recv_pandas(copy=False)
|
||||
|
||||
gevent.joinall([
|
||||
gevent.spawn(send),
|
||||
gevent.spawn(recv)
|
||||
])
|
||||
Reference in New Issue
Block a user