diff --git a/tests/test_assets.py b/tests/test_assets.py index a3509e45..a034be05 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -579,48 +579,9 @@ class AssetFinderTestCase(TestCase): self.assertTrue(issubclass(warning.category, DeprecationWarning)) - def test_lookup_future_in_chain(self): - metadata = { - 2: { - 'symbol': 'ADN15', - 'root_symbol': 'AD', - 'asset_type': 'future', - 'expiration_date': pd.Timestamp('2015-06-15', tz='UTC'), - 'start_date': pd.Timestamp('2015-01-01', tz='UTC') - }, - 1: { - 'symbol': 'ADV15', - 'root_symbol': 'AD', - 'asset_type': 'future', - 'expiration_date': pd.Timestamp('2015-09-14', tz='UTC'), - 'start_date': pd.Timestamp('2015-01-01', tz='UTC') - }, - 0: { - 'symbol': 'ADF16', - 'root_symbol': 'AD', - 'asset_type': 'future', - 'expiration_date': pd.Timestamp('2015-12-14', tz='UTC'), - 'start_date': pd.Timestamp('2015-01-01', tz='UTC') - }, - - } - - finder = AssetFinder(metadata=metadata) - dt = pd.Timestamp('2015-06-19', tz='UTC') - - # Check that the primary and secondary contracts are as expected - primary = finder.lookup_future_in_chain('AD', dt, 0) - secondary = finder.lookup_future_in_chain('AD', dt, 1) - self.assertEqual(primary.sid, 1) - self.assertEqual(secondary.sid, 0) - - # Check that we get None for an invalid contract num - self.assertIsNone(finder.lookup_future_in_chain('AD', dt, -10)) - self.assertIsNone(finder.lookup_future_in_chain('AD', dt, 10)) - self.assertIsNone(finder.lookup_future_in_chain('CL', dt, 1)) - def test_lookup_future_chain(self): metadata = { + # Expires today, so should be valid 2: { 'symbol': 'ADN15', 'root_symbol': 'AD', @@ -641,7 +602,7 @@ class AssetFinderTestCase(TestCase): 'root_symbol': 'AD', 'asset_type': 'future', 'expiration_date': pd.Timestamp('2015-12-14', tz='UTC'), - 'start_date': pd.Timestamp('2015-06-19', tz='UTC') + 'start_date': pd.Timestamp('2015-06-15', tz='UTC') }, # Copy of the above future, but starts trading in August, # so it isn't valid. @@ -656,71 +617,27 @@ class AssetFinderTestCase(TestCase): } finder = AssetFinder(metadata=metadata) - dt = pd.Timestamp('2015-06-19', tz='UTC') + dt = pd.Timestamp('2015-06-15', tz='UTC') + last_year = pd.Timestamp('2014-01-01', tz='UTC') + first_day = pd.Timestamp('2015-01-01', tz='UTC') - # Check that we get the expected number of contract, in the + # Check that we get the expected number of contracts, in the # right order - ad_contracts = finder.lookup_future_chain('AD', dt) + ad_contracts = finder.lookup_future_chain('AD', dt, dt) + self.assertEqual(len(ad_contracts), 3) + self.assertEqual(ad_contracts[0].sid, 2) + self.assertEqual(ad_contracts[1].sid, 1) + self.assertEqual(ad_contracts[2].sid, 0) + + # 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) + + # Check that we get things that start on the knowledge date + ad_contracts = finder.lookup_future_chain('AD', dt, first_day) self.assertEqual(len(ad_contracts), 2) - self.assertEqual(ad_contracts[0].sid, 1) - self.assertEqual(ad_contracts[1].sid, 0) - - def test_lookup_future_by_expiration(self): - metadata = { - 2: { - 'symbol': 'ADN15', - 'root_symbol': 'AD', - 'asset_type': 'future', - 'expiration_date': pd.Timestamp('2015-06-15', tz='UTC'), - 'start_date': pd.Timestamp('2015-01-01', tz='UTC') - }, - 1: { - 'symbol': 'ADV15', - 'root_symbol': 'AD', - 'asset_type': 'future', - 'expiration_date': pd.Timestamp('2015-09-14', tz='UTC'), - 'start_date': pd.Timestamp('2015-01-01', tz='UTC') - }, - 0: { - 'symbol': 'ADF16', - 'root_symbol': 'AD', - 'asset_type': 'future', - 'expiration_date': pd.Timestamp('2015-12-14', tz='UTC'), - 'start_date': pd.Timestamp('2015-01-01', tz='UTC') - }, - - } - - finder = AssetFinder(metadata=metadata) - dt = pd.Timestamp('2015-06-19', tz='UTC') - - # First-of-the-month timestamps - may_15 = pd.Timestamp('2015-05-01', tz='UTC') - june_15 = pd.Timestamp('2015-06-01', tz='UTC') - sept_15 = pd.Timestamp('2015-09-01', tz='UTC') - dec_15 = pd.Timestamp('2015-12-01', tz='UTC') - jan_16 = pd.Timestamp('2016-01-01', tz='UTC') - - # ADV15 is the next valid contract, so check that we get it - # for every ref_date before 9/14/15 - contract = finder.lookup_future_by_expiration('AD', dt, may_15) - self.assertEqual(contract.sid, 1) - - contract = finder.lookup_future_by_expiration('AD', dt, june_15) - self.assertEqual(contract.sid, 1) - - contract = finder.lookup_future_by_expiration('AD', dt, sept_15) - self.assertEqual(contract.sid, 1) - - # ADF16 has the next expiration date after 12/1/15 - contract = finder.lookup_future_by_expiration('AD', dt, dec_15) - self.assertEqual(contract.sid, 0) - - # No contracts exist after 12/14/2015, so we should get none - self.assertIsNone(finder.lookup_future_by_expiration('AD', dt, jan_16)) def test_map_identifier_index_to_sids(self): - # Build an empty finder and some Assets dt = pd.Timestamp('2014-01-01', tz='UTC') finder = AssetFinder() diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 82cd4769..158bbe38 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -29,6 +29,7 @@ from zipline.errors import ( ConsumeAssetMetaDataError, InvalidAssetType, MultipleSymbolsFound, + RootSymbolNotFound, SidAssignmentError, SidNotFound, SymbolNotFound, @@ -231,19 +232,7 @@ class AssetFinder(object): for root_symbol in self.future_chains_cache: self.future_chains_cache[root_symbol].sort(key=exp_key) - def _valid_contracts(self, root_symbol, as_of_date): - """ Returns a list of the currently valid futures contracts - for a given root symbol, sorted by expiration date (the - contracts are sorted when the AssetFinder is built). - """ - try: - return [c for c in self.future_chains_cache[root_symbol] - if c.expiration_date and (c.expiration_date > as_of_date) - and c.start_date and (c.start_date <= as_of_date)] - except KeyError: - return None - - def lookup_future_chain(self, root_symbol, as_of_date): + def lookup_future_chain(self, root_symbol, as_of_date, knowledge_date): """ Return the futures chain for a given root symbol. Parameters @@ -251,77 +240,24 @@ class AssetFinder(object): root_symbol : str Root symbol of the desired future. as_of_date : pd.Timestamp - Date at the time of the lookup. + Date at which the chain determination is rooted. I.e. the + existing contract that expires first after (or on) this date is + the primary contract, etc. + knowledge_date : pd.Timestamp + 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. Returns ------- [Future] """ - root_symbol.upper() - as_of_date = normalize_date(as_of_date) - return self._valid_contracts(root_symbol, as_of_date) - - def lookup_future_in_chain(self, root_symbol, as_of_date, contract_num=0): - """ Find a specific contract in the futures chain for a given - root symbol. - - Parameters - ---------- - root_symbol : str - Root symbol of the desired future. - as_of_date : pd.Timestamp - Date at the time of the lookup. - contract_num : int - 1 for the primary contract, 2 for the secondary, etc., - relative to as_of_date. - - Returns - ------- - Future - The (contract_num)th contract in the futures chain. If none - exits, returns None. - """ - root_symbol.upper() - as_of_date = normalize_date(as_of_date) - - valid_contracts = self._valid_contracts(root_symbol, as_of_date) - - if valid_contracts and contract_num >= 0: - try: - return valid_contracts[contract_num] - except IndexError: - pass - - return None - - def lookup_future_by_expiration(self, root_symbol, as_of_date, ref_date): - """ Find a specific contract in the futures chain by expiration - date. - - Parameters - ---------- - root_symbol : str - Root symbol of the desired future. - as_of_date : pd.Timestamp - Date at the time of the lookup. - ref_date : pd.Timestamp - Reference point for expiration dates. - - Returns - ------- - Future - The valid contract the has the closest expiration date - after ref_date. If none exists, returns None. - """ - root_symbol.upper() - as_of_date = normalize_date(as_of_date) - ref_date = normalize_date(ref_date) - - valid_contracts = self._valid_contracts(root_symbol, as_of_date) - - contracts_after_date = (c for c in valid_contracts - if c.expiration_date > ref_date) - return next(contracts_after_date, None) + try: + return [c for c in self.future_chains_cache[root_symbol] + if c.expiration_date and (as_of_date <= c.expiration_date) + and c.start_date and (c.start_date <= knowledge_date)] + except KeyError: + raise RootSymbolNotFound(root_symbol=root_symbol) def populate_cache(self): """ diff --git a/zipline/errors.py b/zipline/errors.py index 1cb99e8c..54128814 100644 --- a/zipline/errors.py +++ b/zipline/errors.py @@ -215,6 +215,15 @@ Symbol '{symbol}' was not found. """.strip() +class RootSymbolNotFound(ZiplineError): + """ + Raised when a lookup_future_chain() call contains a non-existant symbol. + """ + msg = """ +Root symbol '{root_symbol}' was not found. +""".strip() + + class SidNotFound(ZiplineError): """ Raised when a retrieve_asset() call contains a non-existent sid.