diff --git a/tests/pipeline/test_factor.py b/tests/pipeline/test_factor.py index 3a7d4819..a884049b 100644 --- a/tests/pipeline/test_factor.py +++ b/tests/pipeline/test_factor.py @@ -4,6 +4,7 @@ Tests for Factor terms. from functools import partial from itertools import product from nose_parameterized import parameterized +from unittest import TestCase from toolz import compose from numpy import ( @@ -822,3 +823,17 @@ class FactorTestCase(BasePipelineTestCase): self.assertIs(f.deciles(), f.quantiles(bins=10)) self.assertIs(f.deciles(mask=m), f.quantiles(bins=10, mask=m)) self.assertIsNot(f.deciles(), f.deciles(mask=m)) + + +class ShortReprTestCase(TestCase): + """ + Tests for short_repr methods of Factors. + """ + + def test_demean(self): + r = F().demean().short_repr() + self.assertEqual(r, "GroupedRowTransform('demean')") + + def test_zscore(self): + r = F().zscore().short_repr() + self.assertEqual(r, "GroupedRowTransform('zscore')") diff --git a/zipline/pipeline/classifiers/classifier.py b/zipline/pipeline/classifiers/classifier.py index 7a584eaa..fcd2accd 100644 --- a/zipline/pipeline/classifiers/classifier.py +++ b/zipline/pipeline/classifiers/classifier.py @@ -78,6 +78,9 @@ class Quantiles(SingleInputMixin, Classifier): result[isnan(result)] = self.missing_value return result.astype(int64_dtype) + def short_repr(self): + return type(self).__name__ + '(%d)' % self.params['bins'] + class CustomClassifier(PositiveWindowLengthMixin, CustomTermMixin, Classifier): """ diff --git a/zipline/pipeline/factors/factor.py b/zipline/pipeline/factors/factor.py index 27e0e476..3310523b 100644 --- a/zipline/pipeline/factors/factor.py +++ b/zipline/pipeline/factors/factor.py @@ -576,8 +576,13 @@ class Factor(RestrictedDTypeMixin, ComputableTerm): -------- :meth:`pandas.DataFrame.groupby` """ + # This is a named function so that it has a __name__ for use in the + # graph repr of GroupedRowTransform. + def demean(row): + return row - nanmean(row) + return GroupedRowTransform( - transform=lambda row: row - nanmean(row), + transform=demean, factor=self, mask=mask, groupby=groupby, @@ -637,8 +642,13 @@ class Factor(RestrictedDTypeMixin, ComputableTerm): -------- :meth:`pandas.DataFrame.groupby` """ + # This is a named function so that it has a __name__ for use in the + # graph repr of GroupedRowTransform. + def zscore(row): + return (row - nanmean(row)) / nanstd(row) + return GroupedRowTransform( - transform=lambda row: (row - nanmean(row)) / nanstd(row), + transform=zscore, factor=self, mask=mask, groupby=groupby, @@ -1022,6 +1032,13 @@ class GroupedRowTransform(Factor): self.missing_value, ) + @property + def transform_name(self): + return self._transform.__name__ + + def short_repr(self): + return type(self).__name__ + '(%r)' % self.transform_name + class Rank(SingleInputMixin, Factor): """