MAINT: Use six for Python 3 compatible names and behavior.

Use the six module to import functions and types that are
consistent between Python 2 and 3, so that one code base can
support both versions.

- Use integer types instead of int and long.
- Use string_types instead of basestring.
- Account for iteritems, itervalues, iterkeys.
- Use six.moves for filter and zip, reduce
- Use compatible bytes for md5 hasher.
- xrange and range
This commit is contained in:
Eddie Hebert
2014-01-06 12:21:18 -05:00
parent 40c8c38257
commit b4959e46cf
21 changed files with 95 additions and 47 deletions
+5 -3
View File
@@ -28,6 +28,8 @@ import numpy as np
from nose.tools import timed
from six.moves import range
import zipline.protocol
from zipline.protocol import Event, DATASOURCE_TYPE
@@ -314,7 +316,7 @@ class FinanceTestCase(TestCase):
alternator = 1
order_date = start_date
for i in xrange(order_count):
for i in range(order_count):
blotter.set_date(order_date)
blotter.order(sid, order_amount * alternator ** i, None, None)
@@ -334,7 +336,7 @@ class FinanceTestCase(TestCase):
order_list = oo[sid]
self.assertEqual(order_count, len(order_list))
for i in xrange(order_count):
for i in range(order_count):
order = order_list[i]
self.assertEqual(order.sid, sid)
self.assertEqual(order.amount, order_amount * alternator ** i)
@@ -372,7 +374,7 @@ class FinanceTestCase(TestCase):
self.assertEqual(len(transactions), len(order_list))
total_volume = 0
for i in xrange(len(transactions)):
for i in range(len(transactions)):
txn = transactions[i]
total_volume += txn.amount
if complete_fill:
+3 -1
View File
@@ -24,6 +24,8 @@ import datetime
import pytz
import itertools
from six.moves import range
import zipline.utils.factory as factory
import zipline.finance.performance as perf
from zipline.finance.slippage import Transaction, create_transaction
@@ -431,7 +433,7 @@ class TestDividendPerformance(unittest.TestCase):
pay_date = self.sim_params.first_open
# find pay date that is much later.
for i in xrange(30):
for i in range(30):
pay_date = factory.get_next_trading_dt(pay_date, oneday)
dividend = factory.create_dividend(
1,
+3 -1
View File
@@ -16,6 +16,8 @@ import pandas as pd
import pytz
from itertools import cycle
from six import integer_types
from unittest import TestCase
import zipline.utils.factory as factory
@@ -71,5 +73,5 @@ class TestDataFrameSource(TestCase):
for event in source:
for check_field in check_fields:
self.assertIn(check_field, event)
self.assertTrue(isinstance(event['volume'], (int, long)))
self.assertTrue(isinstance(event['volume'], (integer_types)))
self.assertEqual(stocks_iter.next(), event['sid'])
+4 -2
View File
@@ -20,6 +20,8 @@ import pandas as pd
from datetime import timedelta, datetime
from unittest import TestCase
from six.moves import range
from zipline.utils.test_utils import setup_logger
from zipline.protocol import Event
@@ -64,7 +66,7 @@ class TestEventWindow(TestCase):
self.monday = datetime(2012, 7, 9, 16, tzinfo=pytz.utc)
self.eleven_normal_days = [self.monday + i * timedelta(days=1)
for i in xrange(11)]
for i in range(11)]
# Modify the end of the period slightly to exercise the
# incomplete day logic.
@@ -75,7 +77,7 @@ class TestEventWindow(TestCase):
# Second set of dates to test holiday handling.
self.jul4_monday = datetime(2012, 7, 2, 16, tzinfo=pytz.utc)
self.week_of_jul4 = [self.jul4_monday + i * timedelta(days=1)
for i in xrange(5)]
for i in range(5)]
def test_market_aware_window_normal_week(self):
window = NoopEventWindow(
+5 -3
View File
@@ -20,7 +20,9 @@ import numpy as np
from datetime import datetime
from itertools import groupby, ifilter
from itertools import groupby
from six.moves import filter
from six import iteritems
from operator import attrgetter
from zipline.errors import (
@@ -194,7 +196,7 @@ class TradingAlgorithm(object):
date_sorted = date_sorted_sources(*self.sources)
if source_filter:
date_sorted = ifilter(source_filter, date_sorted)
date_sorted = filter(source_filter, date_sorted)
with_tnfms = sequential_transforms(date_sorted,
*self.transforms)
@@ -305,7 +307,7 @@ class TradingAlgorithm(object):
# Create transforms by wrapping them into StatefulTransforms
self.transforms = []
for namestring, trans_descr in self.registered_transforms.iteritems():
for namestring, trans_descr in iteritems(self.registered_transforms):
sf = StatefulTransform(
trans_descr['class'],
*trans_descr['args'],
+3 -1
View File
@@ -23,6 +23,8 @@ from functools import partial
import requests
import pandas as pd
from six import iteritems
from . loader_utils import (
date_conversion,
source_to_records,
@@ -50,7 +52,7 @@ _BENCHMARK_MAPPING = {
def benchmark_mappings():
return {key: Mapping(*value)
for key, value
in _BENCHMARK_MAPPING.iteritems()}
in iteritems(_BENCHMARK_MAPPING)}
def get_raw_benchmark_data(start_date, end_date, symbol):
+5 -3
View File
@@ -26,6 +26,8 @@ import pandas as pd
from pandas.io.data import DataReader
import pytz
from six import iteritems
from . import benchmarks
from . benchmarks import get_benchmark_returns
@@ -239,7 +241,7 @@ Fetching data from {0}
fp_tr.close()
tr_curves = OrderedDict(sorted(
((dt, c) for dt, c in tr_curves.iteritems()),
((dt, c) for dt, c in iteritems(tr_curves)),
key=lambda t: t[0]))
return benchmark_returns, tr_curves
@@ -291,7 +293,7 @@ must specify stocks or indexes"""
data[stock] = stkd
if indexes is not None:
for name, ticker in indexes.iteritems():
for name, ticker in iteritems(indexes):
print(name)
stkd = DataReader(ticker, 'yahoo', start, end).sort_index()
data[name] = stkd
@@ -327,7 +329,7 @@ def load_from_yahoo(indexes=None,
close_key = 'Adj Close'
else:
close_key = 'Close'
df = pd.DataFrame({key: d[close_key] for key, d in data.iteritems()})
df = pd.DataFrame({key: d[close_key] for key, d in iteritems(data)})
df.index = df.index.tz_localize(pytz.utc)
return df
+3 -1
View File
@@ -30,6 +30,8 @@ from collections import namedtuple
from functools import partial
from six import iteritems
def get_utc_from_exchange_time(naive):
local = pytz.timezone('US/Eastern')
@@ -126,7 +128,7 @@ def _row_cb(mapping, row):
return {
target: apply_mapping(mapping, row)
for target, mapping
in mapping.iteritems()
in iteritems(mapping)
}
+3 -1
View File
@@ -21,6 +21,8 @@ import requests
from collections import OrderedDict
import xml.etree.ElementTree as ET
from six import iteritems
from . loader_utils import (
guarded_conversion,
safe_int,
@@ -61,7 +63,7 @@ _CURVE_MAPPINGS = {
def treasury_mappings(mappings):
return {key: Mapping(*value)
for key, value
in mappings.iteritems()}
in iteritems(mappings)}
class iter_to_stream(object):
+8 -6
View File
@@ -77,6 +77,8 @@ import numpy as np
import pandas as pd
from collections import OrderedDict, defaultdict
from six import iteritems, itervalues
import zipline.protocol as zp
from . position import positiondict
@@ -164,7 +166,7 @@ class PerformancePeriod(object):
payment has been disbursed.
"""
cash_payments = 0.0
for sid, pos in self.positions.iteritems():
for sid, pos in iteritems(self.positions):
cash_payments += pos.update_dividends(todays_date)
# credit our cash balance with the dividend payments, or
@@ -307,7 +309,7 @@ class PerformancePeriod(object):
else:
transactions = \
[y.to_dict()
for x in self.processed_transactions.itervalues()
for x in itervalues(self.processed_transactions)
for y in x]
rval['transactions'] = transactions
@@ -315,9 +317,9 @@ class PerformancePeriod(object):
if dt:
# only include orders modified as of the given dt.
orders = [x.to_dict()
for x in self.orders_by_modified[dt].itervalues()]
for x in itervalues(self.orders_by_modified[dt])]
else:
orders = [x.to_dict() for x in self.orders_by_id.itervalues()]
orders = [x.to_dict() for x in itervalues(self.orders_by_id)]
rval['orders'] = orders
return rval
@@ -352,7 +354,7 @@ class PerformancePeriod(object):
positions = self._positions_store
for sid, pos in self.positions.iteritems():
for sid, pos in iteritems(self.positions):
if sid not in positions:
positions[sid] = zp.Position(sid)
position = positions[sid]
@@ -364,7 +366,7 @@ class PerformancePeriod(object):
def get_positions_list(self):
positions = []
for sid, pos in self.positions.iteritems():
for sid, pos in iteritems(self.positions):
if pos.amount != 0:
positions.append(pos.to_dict())
return positions
+3 -1
View File
@@ -24,6 +24,8 @@ import zipline.utils.math_utils as zp_math
import pandas as pd
from pandas.tseries.tools import normalize_date
from six import iteritems
from . risk import (
alpha,
check_entry,
@@ -359,7 +361,7 @@ algorithm_returns ({algo_count}) in range {start} : {end} on {dt}"
return {k: None
if check_entry(k, v)
else v for k, v in rval.iteritems()}
else v for k, v in iteritems(rval)}
def __repr__(self):
statements = []
+3 -1
View File
@@ -20,6 +20,8 @@ import math
import numpy as np
import numpy.linalg as la
from six import iteritems
from zipline.finance import trading
import pandas as pd
@@ -131,7 +133,7 @@ class RiskMetricsPeriod(object):
}
return {k: None if check_entry(k, v) else v
for k, v in rval.iteritems()}
for k, v in iteritems(rval)}
def __repr__(self):
statements = []
+2
View File
@@ -15,6 +15,8 @@
import heapq
from six.moves import reduce
def _decorate_source(source):
for message in source:
+4 -2
View File
@@ -21,16 +21,18 @@ from hashlib import md5
from datetime import datetime
from zipline.protocol import DATASOURCE_TYPE
from six import iteritems, b
def hash_args(*args, **kwargs):
"""Define a unique string for any set of representable args."""
arg_string = '_'.join([str(arg) for arg in args])
kwarg_string = '_'.join([str(key) + '=' + str(value)
for key, value in kwargs.iteritems()])
for key, value in iteritems(kwargs)])
combined = ':'.join([arg_string, kwarg_string])
hasher = md5()
hasher.update(combined)
hasher.update(b(combined))
return hasher.hexdigest()
+7 -5
View File
@@ -13,6 +13,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from six import iteritems, iterkeys
from . utils.protocol_utils import Enum
# Datasource type should completely determine the other fields of a
@@ -171,7 +173,7 @@ class BarData(object):
del self._data[name]
def __iter__(self):
for sid, data in self._data.iteritems():
for sid, data in iteritems(self._data):
# Allow contains override to filter out sids.
if sid in self:
if len(data):
@@ -179,25 +181,25 @@ class BarData(object):
def iterkeys(self):
# Allow contains override to filter out sids.
return (sid for sid in self._data.iterkeys() if sid in self)
return (sid for sid in iterkeys(self._data) if sid in self)
def keys(self):
# Allow contains override to filter out sids.
return list(self.iterkeys())
def itervalues(self):
return (value for sid, value in self.iteritems())
return (value for sid, value in iteritems(self))
def values(self):
return list(self.itervalues())
def iteritems(self):
return ((sid, value) for sid, value
in self._data.iteritems()
in iteritems(self._data)
if sid in self)
def items(self):
return list(self.iteritems())
return list(iteritems(self))
def __len__(self):
return len(self.keys())
+10 -7
View File
@@ -19,10 +19,13 @@ A source to be used in testing.
import pytz
from itertools import cycle, ifilter, izip
from itertools import cycle
from six.moves import filter, zip
from datetime import datetime, timedelta
import numpy as np
from six.moves import range
from zipline.protocol import (
Event,
DATASOURCE_TYPE
@@ -68,9 +71,9 @@ def date_gen(start=datetime(2006, 6, 6, 12, tzinfo=pytz.utc),
# during trading hours.
# NB: Being inside of trading hours is currently dependent upon the
# count parameter being less than the number of trading minutes in a day
for i in xrange(count):
for i in range(count):
if repeats:
for j in xrange(repeats):
for j in range(repeats):
yield cur
else:
yield cur
@@ -90,7 +93,7 @@ def mock_prices(count):
Utility to generate a stream of mock prices. By default
cycles through values from 0.0 to 10.0, n times.
"""
return (float(i % 10) + 1.0 for i in xrange(count))
return (float(i % 10) + 1.0 for i in range(count))
def mock_volumes(count):
@@ -98,7 +101,7 @@ def mock_volumes(count):
Utility to generate a set of volumes. By default cycles
through values from 100 to 1000, incrementing by 50.
"""
return ((i * 50) % 900 + 100 for i in xrange(count))
return ((i * 50) % 900 + 100 for i in range(count))
class SpecificEquityTrades(object):
@@ -204,7 +207,7 @@ class SpecificEquityTrades(object):
sids = cycle(self.sids)
# Combine the iterators into a single iterator of arguments
arg_gen = izip(sids, prices, volumes, dates)
arg_gen = zip(sids, prices, volumes, dates)
# Convert argument packages into events.
unfiltered = (create_trade(*args, source_id=self.get_hash())
@@ -213,7 +216,7 @@ class SpecificEquityTrades(object):
# If we specified a sid filter, filter out elements that don't
# match the filter.
if self.filter:
filtered = ifilter(
filtered = filter(
lambda event: event.sid in self.filter, unfiltered)
# Otherwise just use all events.
+3 -1
View File
@@ -74,6 +74,8 @@ The algorithm must expose methods:
from copy import deepcopy
import numpy as np
from six.moves import range
from zipline.algorithm import TradingAlgorithm
from zipline.finance.slippage import FixedSlippage
@@ -191,7 +193,7 @@ class TooMuchProcessingAlgorithm(TradingAlgorithm):
def handle_data(self, data):
# Unless we're running on some sort of
# supercomputer this will hit timeout.
for i in xrange(1000000000):
for i in range(1000000000):
self.foo = i
+11 -5
View File
@@ -26,6 +26,12 @@ from numbers import Integral
import pandas as pd
from six import (
string_types,
itervalues,
iteritems
)
from zipline.utils.data import RollingPanel
from zipline.protocol import Event
@@ -187,7 +193,7 @@ class BatchTransform(object):
# enter the batch transform's window IFF a sid filter is not
# specified.
if sids is not None:
if isinstance(sids, (basestring, Integral)):
if isinstance(sids, (string_types, Integral)):
self.static_sids = set([sids])
else:
self.static_sids = set(sids)
@@ -195,7 +201,7 @@ class BatchTransform(object):
self.static_sids = None
self.initial_field_names = fields
if isinstance(self.initial_field_names, basestring):
if isinstance(self.initial_field_names, string_types):
self.initial_field_names = [self.initial_field_names]
self.field_names = set()
@@ -230,7 +236,7 @@ class BatchTransform(object):
Point of entry. Process an event frame.
"""
# extract dates
dts = [event.datetime for event in data._data.itervalues()]
dts = [event.datetime for event in itervalues(data._data)]
# we have to provide the event with a dt. This is only for
# checking if the event is outside the window or not so a
# couple of seconds shouldn't matter. We don't add it to
@@ -238,7 +244,7 @@ class BatchTransform(object):
# sid keys.
event = Event()
event.dt = max(dts)
event.data = {k: v.__dict__ for k, v in data._data.iteritems()
event.data = {k: v.__dict__ for k, v in iteritems(data._data)
# Need to check if data has a 'length' to filter
# out sids without trade data available.
# TODO: expose more of 'no trade available'
@@ -419,7 +425,7 @@ class BatchTransform(object):
# extract field names from sids (price, volume etc), make sure
# every sid has the same fields.
sid_keys = []
for sid in event.data.itervalues():
for sid in itervalues(event.data):
keys = set([name for name, value in sid.items()
if isinstance(value,
(int,
+3 -1
View File
@@ -15,6 +15,8 @@
from collections import defaultdict
from six import string_types
from zipline.transforms.utils import EventWindow, TransformMeta
from zipline.errors import WrongDataForTransform
@@ -31,7 +33,7 @@ class MovingAverage(object):
def __init__(self, fields='price',
market_aware=True, window_length=None, delta=None):
if isinstance(fields, basestring):
if isinstance(fields, string_types):
fields = [fields]
self.fields = fields
+4 -1
View File
@@ -19,6 +19,9 @@ import numpy as np
import pandas as pd
import talib
import copy
from six import iteritems
from zipline.transforms import BatchTransform
@@ -45,7 +48,7 @@ def zipline_wrapper(talib_fn, key_map, data):
for sid in data.minor_axis:
# build talib_data from zipline data
talib_data = dict()
for talib_key, zipline_key in key_map.iteritems():
for talib_key, zipline_key in iteritems(key_map):
# if zipline_key is found, add it to talib_data
if zipline_key in data:
values = data[zipline_key][sid].values
+3 -1
View File
@@ -1,6 +1,8 @@
from logbook import FileHandler
from zipline.finance.blotter import ORDER_STATUS
from six import itervalues
def setup_logger(test, path='test.log'):
test.log_handler = FileHandler(path)
@@ -57,7 +59,7 @@ def assert_single_position(test, zipline):
for order in update['daily_perf']['orders']:
orders_by_id[order['id']] = order
for order in orders_by_id.itervalues():
for order in itervalues(orders_by_id):
test.assertEqual(
order['status'],
ORDER_STATUS.FILLED,