diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index fced35a9..948baa05 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -14,6 +14,7 @@ # limitations under the License. import numpy as np import pandas as pd +from pandas.util.testing import assert_series_equal from zipline.data.data_portal import DataPortal from zipline.errors import ( @@ -117,6 +118,12 @@ class TestBenchmark(WithDataPortal, WithSimParams, WithTradingCalendars, manually_calculated[idx + 1] ) + # compare a slice of the data + assert_series_equal( + source.get_range(days_to_use[1], days_to_use[10]), + manually_calculated[1:11] + ) + def test_asset_not_trading(self): benchmark = self.env.asset_finder.retrieve_asset(3) benchmark_start = benchmark.start_date diff --git a/zipline/sources/benchmark_source.py b/zipline/sources/benchmark_source.py index c20892f4..309ca8d2 100644 --- a/zipline/sources/benchmark_source.py +++ b/zipline/sources/benchmark_source.py @@ -72,6 +72,9 @@ class BenchmarkSource(object): def get_value(self, dt): return self._precalculated_series.loc[dt] + def get_range(self, start_dt, end_dt): + return self._precalculated_series.loc[start_dt:end_dt] + def _validate_benchmark(self, benchmark_asset): # check if this security has a stock dividend. if so, raise an # error suggesting that the user pick a different asset to use