diff --git a/tests/test_assets.py b/tests/test_assets.py index f70ac69a..98b0e06e 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -587,6 +587,10 @@ class AssetFinderTestCase(TestCase): self.assertEqual(ad_contracts[0].sid, 1) self.assertEqual(ad_contracts[1].sid, 0) + # Check that pd.NaT for knowledge_date uses the value of as_of_date + ad_contracts = finder.lookup_future_chain('AD', dt, pd.NaT) + self.assertEqual(len(ad_contracts), 2) + # Check that we get nothing if our knowledge date is last year ad_contracts = finder.lookup_future_chain('AD', dt, last_year) self.assertEqual(len(ad_contracts), 0) @@ -595,6 +599,10 @@ class AssetFinderTestCase(TestCase): ad_contracts = finder.lookup_future_chain('AD', dt, first_day) self.assertEqual(len(ad_contracts), 1) + # Check that pd.NaT for as_of_date gives the whole chain + ad_contracts = finder.lookup_future_chain('AD', pd.NaT, first_day) + self.assertEqual(len(ad_contracts), 4) + def test_map_identifier_index_to_sids(self): # Build an empty finder and some Assets dt = pd.Timestamp('2014-01-01', tz='UTC') diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index c847daf6..c012996d 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -440,14 +440,18 @@ class AssetFinder(object): ---------- root_symbol : str Root symbol of the desired future. - as_of_date : pd.Timestamp + as_of_date : pd.Timestamp or pd.NaT Date at which the chain determination is rooted. I.e. the existing contract whose notice date is first after this - date is the primary contract, etc. - knowledge_date : pd.Timestamp + date is the primary contract, etc. If NaT is given, the + chain is unbounded, and all contracts for this root symbol + are returned. + knowledge_date : pd.Timestamp or pd.NaT Date for determining which contracts exist for inclusion in this chain. Contracts exist only if they have a start_date - on or before this date. + on or before this date. If NaT is given and as_of_date is + is not NaT, the value of as_of_date is used for + knowledge_date. Returns ------- @@ -462,16 +466,34 @@ class AssetFinder(object): root symbol. """ c = self.conn.cursor() - t = {'root_symbol': root_symbol, - 'as_of_date': as_of_date.value, - 'knowledge_date': knowledge_date.value} - c.execute(""" - select sid from futures - where root_symbol=:root_symbol - and :as_of_date < notice_date - and start_date <= :knowledge_date - order by notice_date asc - """, t) + + if as_of_date is pd.NaT: + # If the as_of_date is NaT, get all contracts for this + # root symbol. + t = {'root_symbol': root_symbol} + c.execute(""" + select sid from futures + where root_symbol=:root_symbol + order by notice_date asc + """, t) + else: + if knowledge_date is pd.NaT: + # If knowledge_date is NaT, default to using as_of_date + t = {'root_symbol': root_symbol, + 'as_of_date': as_of_date.value, + 'knowledge_date': as_of_date.value} + else: + t = {'root_symbol': root_symbol, + 'as_of_date': as_of_date.value, + 'knowledge_date': knowledge_date.value} + + c.execute(""" + select sid from futures + where root_symbol=:root_symbol + and :as_of_date < notice_date + and start_date <= :knowledge_date + order by notice_date asc + """, t) sids = [r[0] for r in c.fetchall()] if not sids: # Check if root symbol exists.