def test_listing_currency(self):
        # Test loading on all assets.
        all_assets = np.array(list(self.assets))
        all_results = self.daily_bar_reader.currency_codes(all_assets)
        all_expected = self.make_equity_daily_bar_currency_codes(
            self.DAILY_BARS_TEST_QUERY_COUNTRY_CODE, all_assets,
        ).values
        assert_equal(all_results, all_expected)

        self.assertEqual(all_results.dtype, np.dtype(object))
        for code in all_results:
            self.assertIsInstance(code, str)

        # Check all possible subsets of assets.
        for indices in map(list, powerset(range(len(all_assets)))):
            # Empty queries aren't currently supported.
            if not indices:
                continue
            assets = all_assets[indices]
            results = self.daily_bar_reader.currency_codes(assets)
            expected = all_expected[indices]

            assert_equal(results, expected)
示例#2
0
    def test_compute_lifetimes(self):
        assets_per_exchange = 4
        trading_day = self.trading_calendar.day
        first_start = pd.Timestamp('2015-04-01', tz='UTC')

        equities = pd.concat(
            [
                make_rotating_equity_info(
                    num_assets=assets_per_exchange,
                    first_start=first_start,
                    frequency=trading_day,
                    periods_between_starts=3,
                    asset_lifetime=5,
                    exchange=exchange,
                ) for exchange in (
                    'US_EXCHANGE_1',
                    'US_EXCHANGE_2',
                    'CA_EXCHANGE',
                    'JP_EXCHANGE',
                )
            ],
            ignore_index=True,
        )
        # make every symbol unique
        equities['symbol'] = list(string.ascii_uppercase[:len(equities)])
        equities['real_sid'] = equities['symbol']

        # shuffle up the sids so they are not contiguous per exchange
        sids = np.arange(len(equities))
        np.random.RandomState(1337).shuffle(sids)
        equities.index = sids
        permute_sid = dict(zip(sids, range(len(sids)))).__getitem__

        exchanges = pd.DataFrame.from_records([
            {
                'exchange': 'US_EXCHANGE_1',
                'country_code': 'US'
            },
            {
                'exchange': 'US_EXCHANGE_2',
                'country_code': 'US'
            },
            {
                'exchange': 'CA_EXCHANGE',
                'country_code': 'CA'
            },
            {
                'exchange': 'JP_EXCHANGE',
                'country_code': 'JP'
            },
        ])
        sids_by_country = {
            'US': equities.index[:2 * assets_per_exchange],
            'CA':
            equities.index[2 * assets_per_exchange:3 * assets_per_exchange],
            'JP': equities.index[3 * assets_per_exchange:],
        }

        self.write_assets(equities=equities, exchanges=exchanges)
        finder = self.asset_finder

        all_dates = pd.date_range(
            start=first_start,
            end=equities.end_date.max(),
            freq=trading_day,
        )

        for dates in all_subindices(all_dates):
            expected_with_start_raw = full(
                shape=(len(dates), assets_per_exchange),
                fill_value=False,
                dtype=bool,
            )
            expected_no_start_raw = full(
                shape=(len(dates), assets_per_exchange),
                fill_value=False,
                dtype=bool,
            )

            for i, date in enumerate(dates):
                it = equities.iloc[:4][['start_date',
                                        'end_date']].itertuples(index=False, )
                for j, (start, end) in enumerate(it):
                    # This way of doing the checks is redundant, but very
                    # clear.
                    if start <= date <= end:
                        expected_with_start_raw[i, j] = True
                        if start < date:
                            expected_no_start_raw[i, j] = True

            for country_codes in powerset(exchanges.country_code.unique()):
                expected_sids = pd.Int64Index(
                    sorted(
                        concat(sids_by_country[country_code]
                               for country_code in country_codes)))
                permuted_sids = [
                    sid for sid in sorted(expected_sids, key=permute_sid)
                ]
                tile_count = len(country_codes) + ('US' in country_codes)
                expected_with_start = pd.DataFrame(
                    data=np.tile(
                        expected_with_start_raw,
                        tile_count,
                    ),
                    index=dates,
                    columns=pd.Int64Index(permuted_sids),
                )
                result = finder.lifetimes(
                    dates,
                    include_start_date=True,
                    country_codes=country_codes,
                )
                assert_equal(result.columns, expected_sids)
                result = result[permuted_sids]
                assert_equal(result, expected_with_start)

                expected_no_start = pd.DataFrame(
                    data=np.tile(
                        expected_no_start_raw,
                        tile_count,
                    ),
                    index=dates,
                    columns=pd.Int64Index(permuted_sids),
                )
                result = finder.lifetimes(
                    dates,
                    include_start_date=False,
                    country_codes=country_codes,
                )
                assert_equal(result.columns, expected_sids)
                result = result[permuted_sids]
                assert_equal(result, expected_no_start)