diff --git a/tests/test_assets.py b/tests/test_assets.py index cf52be8d..bb6142cb 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -720,6 +720,43 @@ class AssetFinderTestCase(WithTradingCalendars, ZiplineTestCase): self.assertEqual(results, expected) self.assertEqual(missing, []) + def test_lookup_none_raises(self): + """ + If lookup_symbol is vectorized across multiple symbols, and one of them + is None, want to raise a TypeError. + """ + + with self.assertRaises(TypeError): + self.asset_finder.lookup_symbol(None, pd.Timestamp('2013-01-01')) + + def test_lookup_mult_are_one(self): + """ + Ensure that multiple symbols that return the same sid are collapsed to + a single returned asset. + """ + + date = pd.Timestamp('2013-01-01', tz='UTC') + + df = pd.DataFrame.from_records( + [ + { + 'sid': 1, + 'symbol': symbol, + 'start_date': date.value, + 'end_date': (date + timedelta(days=30)).value, + 'exchange': 'NYSE', + } + for symbol in ('FOOB', 'FOO_B') + ] + ) + self.write_assets(equities=df) + finder = self.asset_finder + + # If we are able to resolve this with any result, means that we did not + # raise a MultipleSymbolError. + result = finder.lookup_symbol('FOO/B', date + timedelta(1), fuzzy=True) + self.assertEqual(result.sid, 1) + def test_lookup_generic_handle_missing(self): data = pd.DataFrame.from_records( [ diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 08c481e3..1e25181b 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -640,21 +640,23 @@ class AssetFinder(object): options=set(options), ) - options = [] + options = {} for start, end, sid, sym in owners: if start <= as_of_date < end: # see which fuzzy symbols were owned on the asof date. - options.append((sid, sym)) + options[sid] = sym if not options: # no equity owned the fuzzy symbol on the date requested raise SymbolNotFound(symbol=symbol) + sid_keys = options.keys() + # If there was only one owner, or there is a fuzzy and non-fuzzy which + # map to the same sid, return it. if len(options) == 1: - # there was only one owner, return it - return self.retrieve_asset(options[0][0]) + return self.retrieve_asset(sid_keys[0]) - for sid, sym in options: + for sid, sym in options.items(): if sym == symbol: # look for an exact match on the asof date return self.retrieve_asset(sid) @@ -663,10 +665,7 @@ class AssetFinder(object): # there are no exact matches raise MultipleSymbolsFound( symbol=symbol, - options=set(map( - compose(self.retrieve_asset, itemgetter(0)), - options, - )), + options=[self.retrieve_asset(s) for s in sid_keys], ) def lookup_symbol(self, symbol, as_of_date, fuzzy=False): @@ -704,8 +703,8 @@ class AssetFinder(object): ``as_of_date``. """ if symbol is None: - raise TypeError("Cannot lookup symbol of type NoneType for " - "as of date %s" % as_of_date) + raise TypeError("Cannot lookup asset for symbol of None for " + "as of date %s." % as_of_date) if fuzzy: return self._lookup_symbol_fuzzy(symbol, as_of_date)