Skip to content

Commit

Permalink
ENH: Enhancements to TradingEnvironment.
Browse files Browse the repository at this point in the history 
Adds a suite of new functions for querying data from the trading calendar.

These include:
      `previous_trading_day`
      `minutes_for_days_in_range` (minutely version of `days_in_range`)
      `previous_open_and_close` (inverse of `next_open_and_close`)
      `next_market_minute`
      `previous_market_minute`
      `open_close_window` (get a range of opens/closes with slicing semantics)
      `market_minute_window` (get a range of minutes with slicing semantics)

Also refactors `test_finance` to move `TradingEnvironment` tests into their own
TestCase.
  • Loading branch information
@ssanderson
ssanderson committed May 30, 2014
1 parent c7d7c1f commit a98b690
Show file tree
Hide file tree
Showing 2 changed files with 302 additions and 86 deletions.
263 changes: 181 additions & 82 deletions tests/test_finance.py
View file
Expand Up @@ -40,6 +40,7 @@
from zipline.gens.composites import date_sorted_sources

from zipline.finance import trading
from zipline.finance.trading import TradingEnvironment
from zipline.finance.execution import MarketOrder, LimitOrder
from zipline.finance.trading import SimulationParameters

Expand Down Expand Up @@ -80,88 +81,6 @@ def test_factory_daily(self):
self.assertTrue(trade.dt > prev.dt)
prev = trade

@timed(DEFAULT_TIMEOUT)
def test_trading_environment(self):
# holidays taken from: http://www.nyse.com/press/1191407641943.html
new_years = datetime(2008, 1, 1, tzinfo=pytz.utc)
mlk_day = datetime(2008, 1, 21, tzinfo=pytz.utc)
presidents = datetime(2008, 2, 18, tzinfo=pytz.utc)
good_friday = datetime(2008, 3, 21, tzinfo=pytz.utc)
memorial_day = datetime(2008, 5, 26, tzinfo=pytz.utc)
july_4th = datetime(2008, 7, 4, tzinfo=pytz.utc)
labor_day = datetime(2008, 9, 1, tzinfo=pytz.utc)
tgiving = datetime(2008, 11, 27, tzinfo=pytz.utc)
christmas = datetime(2008, 5, 25, tzinfo=pytz.utc)
a_saturday = datetime(2008, 8, 2, tzinfo=pytz.utc)
a_sunday = datetime(2008, 10, 12, tzinfo=pytz.utc)
holidays = [
new_years,
mlk_day,
presidents,
good_friday,
memorial_day,
july_4th,
labor_day,
tgiving,
christmas,
a_saturday,
a_sunday
]

for holiday in holidays:
self.assertTrue(not trading.environment.is_trading_day(holiday))

first_trading_day = datetime(2008, 1, 2, tzinfo=pytz.utc)
last_trading_day = datetime(2008, 12, 31, tzinfo=pytz.utc)
workdays = [first_trading_day, last_trading_day]

for workday in workdays:
self.assertTrue(trading.environment.is_trading_day(workday))

def test_simulation_parameters(self):
env = SimulationParameters(
period_start=datetime(2008, 1, 1, tzinfo=pytz.utc),
period_end=datetime(2008, 12, 31, tzinfo=pytz.utc),
capital_base=100000,
)

self.assertTrue(env.last_close.month == 12)
self.assertTrue(env.last_close.day == 31)

@timed(DEFAULT_TIMEOUT)
def test_sim_params_days_in_period(self):

# January 2008
# Su Mo Tu We Th Fr Sa
# 1 2 3 4 5
# 6 7 8 9 10 11 12
# 13 14 15 16 17 18 19
# 20 21 22 23 24 25 26
# 27 28 29 30 31

env = SimulationParameters(
period_start=datetime(2007, 12, 31, tzinfo=pytz.utc),
period_end=datetime(2008, 1, 7, tzinfo=pytz.utc),
capital_base=100000,
)

expected_trading_days = (
datetime(2007, 12, 31, tzinfo=pytz.utc),
# Skip new years
# holidays taken from: http://www.nyse.com/press/1191407641943.html
datetime(2008, 1, 2, tzinfo=pytz.utc),
datetime(2008, 1, 3, tzinfo=pytz.utc),
datetime(2008, 1, 4, tzinfo=pytz.utc),
# Skip Saturday
# Skip Sunday
datetime(2008, 1, 7, tzinfo=pytz.utc)
)

num_expected_trading_days = 5
self.assertEquals(num_expected_trading_days, env.days_in_period)
np.testing.assert_array_equal(expected_trading_days,
env.trading_days.tolist())

@timed(EXTENDED_TIMEOUT)
def test_full_zipline(self):
# provide enough trades to ensure all orders are filled.
Expand Down Expand Up @@ -429,3 +348,183 @@ def test_blotter_processes_splits(self):
self.assertEqual(300, fls_order['amount'])
self.assertEqual(3.33, fls_order['limit'])
self.assertEqual(2, fls_order['sid'])


class TradingEnvironmentTestCase(TestCase):
"""
Tests for date management utilities in zipline.finance.trading.
"""

def setUp(self):
setup_logger(self)

def tearDown(self):
teardown_logger(self)

@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()

@timed(DEFAULT_TIMEOUT)
def test_is_trading_day(self):
# holidays taken from: http://www.nyse.com/press/1191407641943.html
new_years = datetime(2008, 1, 1, tzinfo=pytz.utc)
mlk_day = datetime(2008, 1, 21, tzinfo=pytz.utc)
presidents = datetime(2008, 2, 18, tzinfo=pytz.utc)
good_friday = datetime(2008, 3, 21, tzinfo=pytz.utc)
memorial_day = datetime(2008, 5, 26, tzinfo=pytz.utc)
july_4th = datetime(2008, 7, 4, tzinfo=pytz.utc)
labor_day = datetime(2008, 9, 1, tzinfo=pytz.utc)
tgiving = datetime(2008, 11, 27, tzinfo=pytz.utc)
christmas = datetime(2008, 5, 25, tzinfo=pytz.utc)
a_saturday = datetime(2008, 8, 2, tzinfo=pytz.utc)
a_sunday = datetime(2008, 10, 12, tzinfo=pytz.utc)
holidays = [
new_years,
mlk_day,
presidents,
good_friday,
memorial_day,
july_4th,
labor_day,
tgiving,
christmas,
a_saturday,
a_sunday
]

for holiday in holidays:
self.assertTrue(not self.env.is_trading_day(holiday))

first_trading_day = datetime(2008, 1, 2, tzinfo=pytz.utc)
last_trading_day = datetime(2008, 12, 31, tzinfo=pytz.utc)
workdays = [first_trading_day, last_trading_day]

for workday in workdays:
self.assertTrue(self.env.is_trading_day(workday))

def test_simulation_parameters(self):
env = SimulationParameters(
period_start=datetime(2008, 1, 1, tzinfo=pytz.utc),
period_end=datetime(2008, 12, 31, tzinfo=pytz.utc),
capital_base=100000,
)

self.assertTrue(env.last_close.month == 12)
self.assertTrue(env.last_close.day == 31)

@timed(DEFAULT_TIMEOUT)
def test_sim_params_days_in_period(self):

# January 2008
# Su Mo Tu We Th Fr Sa
# 1 2 3 4 5
# 6 7 8 9 10 11 12
# 13 14 15 16 17 18 19
# 20 21 22 23 24 25 26
# 27 28 29 30 31

env = SimulationParameters(
period_start=datetime(2007, 12, 31, tzinfo=pytz.utc),
period_end=datetime(2008, 1, 7, tzinfo=pytz.utc),
capital_base=100000,
)

expected_trading_days = (
datetime(2007, 12, 31, tzinfo=pytz.utc),
# Skip new years
# holidays taken from: http://www.nyse.com/press/1191407641943.html
datetime(2008, 1, 2, tzinfo=pytz.utc),
datetime(2008, 1, 3, tzinfo=pytz.utc),
datetime(2008, 1, 4, tzinfo=pytz.utc),
# Skip Saturday
# Skip Sunday
datetime(2008, 1, 7, tzinfo=pytz.utc)
)

num_expected_trading_days = 5
self.assertEquals(num_expected_trading_days, env.days_in_period)
np.testing.assert_array_equal(expected_trading_days,
env.trading_days.tolist())

@timed(DEFAULT_TIMEOUT)
def test_market_minute_window(self):

# January 2008
# Su Mo Tu We Th Fr Sa
# 1 2 3 4 5
# 6 7 8 9 10 11 12
# 13 14 15 16 17 18 19
# 20 21 22 23 24 25 26
# 27 28 29 30 31

us_east = pytz.timezone('US/Eastern')
utc = pytz.utc

# 10:01 AM Eastern on January 7th..
start = us_east.localize(datetime(2008, 1, 7, 10, 1))
utc_start = start.astimezone(utc)

# Get the next 10 minutes
minutes = self.env.market_minute_window(
utc_start, 10,
)
self.assertEqual(len(minutes), 10)
for i in range(10):
self.assertEqual(minutes[i], utc_start + timedelta(minutes=i))

# Get the previous 10 minutes.
minutes = self.env.market_minute_window(
utc_start, 10, step=-1,
)
self.assertEqual(len(minutes), 10)
for i in range(10):
self.assertEqual(minutes[i], utc_start + timedelta(minutes=-i))

# Get the next 900 minutes, including utc_start, rolling over into the
# next two days.
# Should include:
# Today: 10:01 AM -> 4:00 PM (360 minutes)
# Tomorrow: 9:31 AM -> 4:00 PM (390 minutes, 750 total)
# Last Day: 9:31 AM -> 12:00 PM (150 minutes, 900 total)
minutes = self.env.market_minute_window(
utc_start, 900,
)
today = self.env.market_minutes_for_day(start)[30:]
tomorrow = self.env.market_minutes_for_day(
start + timedelta(days=1)
)
last_day = self.env.market_minutes_for_day(
start + timedelta(days=2))[:150]

self.assertEqual(len(minutes), 900)
self.assertEqual(minutes[0], utc_start)
self.assertTrue(all(today == minutes[:360]))
self.assertTrue(all(tomorrow == minutes[360:750]))
self.assertTrue(all(last_day == minutes[750:]))

# Get the previous 801 minutes, including utc_start, rolling over into
# Friday the 4th and Thursday the 3rd.
# Should include:
# Today: 10:01 AM -> 9:31 AM (31 minutes)
# Friday: 4:00 PM -> 9:31 AM (390 minutes, 421 total)
# Thursday: 4:00 PM -> 9:41 AM (380 minutes, 801 total)
minutes = self.env.market_minute_window(
utc_start, 801, step=-1,
)

today = self.env.market_minutes_for_day(start)[30::-1]
# minus an extra two days from each of these to account for the two
# weekend days we skipped
friday = self.env.market_minutes_for_day(
start + timedelta(days=-3),
)[::-1]
thursday = self.env.market_minutes_for_day(
start + timedelta(days=-4),
)[:9:-1]

self.assertEqual(len(minutes), 801)
self.assertEqual(minutes[0], utc_start)
self.assertTrue(all(today == minutes[:31]))
self.assertTrue(all(friday == minutes[31:421]))
self.assertTrue(all(thursday == minutes[421:]))

0 comments on commit a98b690

Please sign in to comment.