diff --git a/tests/pipeline/test_engine.py b/tests/pipeline/test_engine.py index 3cef954b..7f748650 100644 --- a/tests/pipeline/test_engine.py +++ b/tests/pipeline/test_engine.py @@ -41,7 +41,9 @@ from scipy.stats.stats import linregress, pearsonr, spearmanr from six import iteritems, itervalues from toolz import merge +from zipline.assets import Equity from zipline.assets.synthetic import make_rotating_equity_info +from zipline.errors import NonExistentAssetInTimeFrame from zipline.lib.adjustment import MULTIPLY from zipline.lib.labelarray import LabelArray from zipline.pipeline import CustomFactor, Pipeline @@ -1431,6 +1433,59 @@ class ParameterizedFactorTestCase(WithTradingEnvironment, ZiplineTestCase): ), ) + def test_correlation_and_regression_with_bad_asset(self): + """ + Test that `RollingPearsonOfReturns`, `RollingSpearmanOfReturns` and + `RollingLinearRegressionOfReturns` raise the proper exception when + given a nonexistent target asset. + """ + start_date_index = 6 + end_date_index = 10 + my_asset = Equity(0) + + # This filter is arbitrary; the important thing is that we test each + # factor both with and without a specified mask. + my_asset_filter = AssetID().eq(1) + + for mask in (NotSpecified, my_asset_filter): + pearson_factor = RollingPearsonOfReturns( + target=my_asset, + returns_length=3, + correlation_length=3, + mask=mask, + ) + spearman_factor = RollingSpearmanOfReturns( + target=my_asset, + returns_length=3, + correlation_length=3, + mask=mask, + ) + regression_factor = RollingLinearRegressionOfReturns( + target=my_asset, + returns_length=3, + regression_length=3, + mask=mask, + ) + + with self.assertRaises(NonExistentAssetInTimeFrame): + self.engine.run_pipeline( + Pipeline(columns={'pearson_factor': pearson_factor}), + self.dates[start_date_index], + self.dates[end_date_index], + ) + with self.assertRaises(NonExistentAssetInTimeFrame): + self.engine.run_pipeline( + Pipeline(columns={'spearman_factor': spearman_factor}), + self.dates[start_date_index], + self.dates[end_date_index], + ) + with self.assertRaises(NonExistentAssetInTimeFrame): + self.engine.run_pipeline( + Pipeline(columns={'regression_factor': regression_factor}), + self.dates[start_date_index], + self.dates[end_date_index], + ) + class StringColumnTestCase(WithSeededRandomPipelineEngine, ZiplineTestCase): diff --git a/zipline/pipeline/factors/technical.py b/zipline/pipeline/factors/technical.py index 88057bc2..88b141a5 100644 --- a/zipline/pipeline/factors/technical.py +++ b/zipline/pipeline/factors/technical.py @@ -26,7 +26,7 @@ from scipy.stats import linregress, pearsonr, spearmanr from zipline.pipeline.data import USEquityPricing from zipline.pipeline.filters import SingleAsset from zipline.pipeline.mixins import SingleInputMixin -from zipline.pipeline.term import NotSpecified +from zipline.pipeline.term import AssetExists, NotSpecified from zipline.utils.numpy_utils import ignore_nanwarnings from zipline.utils.input_validation import expect_types from zipline.utils.math_utils import ( @@ -183,9 +183,12 @@ class _RollingCorrelationOfReturns(CustomFactor, SingleInputMixin): correlation_length, mask=NotSpecified, **kwargs): - if mask is not NotSpecified: - # Make sure we do not filter out the asset of interest. - mask = mask | SingleAsset(asset=target) + if mask is NotSpecified: + mask = AssetExists() + + # Make sure we do not filter out the asset of interest. + mask = mask | SingleAsset(asset=target) + return super(_RollingCorrelationOfReturns, cls).__new__( cls, target=target, @@ -390,9 +393,12 @@ class RollingLinearRegressionOfReturns(CustomFactor, SingleInputMixin): regression_length, mask=NotSpecified, **kwargs): - if mask is not NotSpecified: - # Make sure we do not filter out the asset of interest. - mask = mask | SingleAsset(asset=target) + if mask is NotSpecified: + mask = AssetExists() + + # Make sure we do not filter out the asset of interest. + mask = mask | SingleAsset(asset=target) + return super(RollingLinearRegressionOfReturns, cls).__new__( cls, target=target,