diff --git a/tests/test_assets.py b/tests/test_assets.py index b6fd04b7..6038968e 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -233,15 +233,30 @@ class AssetTestCase(TestCase): class TestFuture(TestCase): - future = Future( - 2468, - symbol='OMH15', - root_symbol='OM', - notice_date=pd.Timestamp('2014-01-20', tz='UTC'), - expiration_date=pd.Timestamp('2014-02-20', tz='UTC'), - auto_close_date=pd.Timestamp('2014-01-18', tz='UTC'), - contract_multiplier=500 - ) + + @classmethod + def setUpClass(cls): + cls.future = Future( + 2468, + symbol='OMH15', + root_symbol='OM', + notice_date=pd.Timestamp('2014-01-20', tz='UTC'), + expiration_date=pd.Timestamp('2014-02-20', tz='UTC'), + auto_close_date=pd.Timestamp('2014-01-18', tz='UTC'), + contract_multiplier=500 + ) + cls.future2 = Future( + 0, + symbol='CLG06', + root_symbol='CL', + start_date=pd.Timestamp('2005-12-01', tz='UTC'), + notice_date=pd.Timestamp('2005-12-20', tz='UTC'), + expiration_date=pd.Timestamp('2006-01-20', tz='UTC') + ) + env = TradingEnvironment() + env.write_data(futures_identifiers=[TestFuture.future, + TestFuture.future2]) + cls.asset_finder = env.asset_finder def test_str(self): strd = self.future.__str__() @@ -280,6 +295,41 @@ class TestFuture(TestCase): def test_root_symbol(self): self.assertEqual('OM', self.future.root_symbol) + def test_lookup_future_symbol(self): + """ + Test the lookup_future_symbol method. + """ + om = TestFuture.asset_finder.lookup_future_symbol('OMH15') + self.assertEqual(om.sid, 2468) + self.assertEqual(om.symbol, 'OMH15') + self.assertEqual(om.root_symbol, 'OM') + self.assertEqual(om.notice_date, pd.Timestamp('2014-01-20', tz='UTC')) + self.assertEqual(om.expiration_date, + pd.Timestamp('2014-02-20', tz='UTC')) + self.assertEqual(om.auto_close_date, + pd.Timestamp('2014-01-18', tz='UTC')) + + cl = TestFuture.asset_finder.lookup_future_symbol('CLG06') + self.assertEqual(cl.sid, 0) + self.assertEqual(cl.symbol, 'CLG06') + self.assertEqual(cl.root_symbol, 'CL') + self.assertEqual(cl.start_date, pd.Timestamp('2005-12-01', tz='UTC')) + self.assertEqual(cl.notice_date, pd.Timestamp('2005-12-20', tz='UTC')) + self.assertEqual(cl.expiration_date, + pd.Timestamp('2006-01-20', tz='UTC')) + + with self.assertRaises(SymbolNotFound): + TestFuture.asset_finder.lookup_future_symbol('') + + with self.assertRaises(SymbolNotFound): + TestFuture.asset_finder.lookup_future_symbol('#&?!') + + with self.assertRaises(SymbolNotFound): + TestFuture.asset_finder.lookup_future_symbol('FOOBAR') + + with self.assertRaises(SymbolNotFound): + TestFuture.asset_finder.lookup_future_symbol('XXX99') + class AssetFinderTestCase(TestCase): @@ -599,10 +649,9 @@ class AssetFinderTestCase(TestCase): 'notice_date': pd.Timestamp('2015-11-16', tz='UTC'), 'start_date': pd.Timestamp('2015-05-14', tz='UTC') }, - # Copy of the above future, but starts trading in August, - # so it isn't valid. + # Starts trading in August, so not valid. 3: { - 'symbol': 'ADF16', + 'symbol': 'ADX16', 'root_symbol': 'AD', 'asset_type': 'future', 'notice_date': pd.Timestamp('2015-11-16', tz='UTC'), diff --git a/zipline/assets/asset_writer.py b/zipline/assets/asset_writer.py index 99ef87fe..5a730360 100644 --- a/zipline/assets/asset_writer.py +++ b/zipline/assets/asset_writer.py @@ -326,7 +326,7 @@ class AssetDBWriter(with_metaclass(ABCMeta)): nullable=False, primary_key=True, ), - sa.Column('symbol', sa.Text), + sa.Column('symbol', sa.Text, unique=True, index=True), sa.Column( 'root_symbol', sa.Text, diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index a7ceb7af..108dc5e3 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -175,6 +175,10 @@ class AssetFinder(object): def _select_asset_by_sid(asset_tbl, sid): return sa.select([asset_tbl]).where(asset_tbl.c.sid == int(sid)) + @staticmethod + def _select_asset_by_symbol(asset_tbl, symbol): + return sa.select([asset_tbl]).where(asset_tbl.c.symbol == symbol) + def _retrieve_asset(self, sid, cache, asset_tbl, asset_type): try: return cache[sid] @@ -297,6 +301,49 @@ class AssetFinder(object): )) ) + def lookup_future_symbol(self, symbol): + """ Return the Future object for a given symbol. + + Parameters + ---------- + symbol : str + The symbol of the desired contract. + + Returns + ------- + Future + A Future object. + + Raises + ------ + SymbolNotFound + Raised when no contract named 'symbol' is found. + + """ + + data = self._select_asset_by_symbol(self.futures_contracts, symbol)\ + .execute().fetchone() + + # If no data found, raise an exception + if not data: + raise SymbolNotFound(symbol=symbol) + + # If we find a contract, check whether it's been cached + try: + return self._future_cache[data['sid']] + except KeyError: + pass + + # Build the Future object from its parameters + data = dict(data.items()) + _convert_asset_timestamp_fields(data) + future = Future(**data) + + # Cache the Future object. + self._future_cache[data['sid']] = future + + return future + def lookup_future_chain(self, root_symbol, as_of_date, knowledge_date): """ Return the futures chain for a given root symbol.