ENH: Add eq and __ne__ to Classifier.

tests/pipeline/test_classifier.py

@@ -1,7 +1,7 @@
 import numpy as np
 
-from zipline.pipeline import Classifier, TermGraph
-from zipline.testing import check_arrays, parameter_space
+from zipline.pipeline import Classifier
+from zipline.testing import parameter_space
 from zipline.utils.numpy_utils import int64_dtype
 
 from .base import BasePipelineTestCase
@@ -18,26 +18,103 @@ class ClassifierTestCase(BasePipelineTestCase):
             inputs = ()
             window_length = 0
 
+        c = C()
+
         # There's no significance to the values here other than that they
         # contain a mix of missing and non-missing values.
         data = np.array([[-1,  1,  0, 2],
                          [3,   0,  1, 0],
                          [-5,  0, -1, 0],
-                         [-3,  1,  2, 2]], dtype=int)
+                         [-3,  1,  2, 2]], dtype=int64_dtype)
 
-        c = C()
-        graph = TermGraph(
-            {
+        self.check_terms(
+            terms={
                 'isnull': c.isnull(),
                 'notnull': c.notnull()
-            }
-        )
-
-        results = self.run_graph(
-            graph,
+            },
+            expected={
+                'isnull': data == mv,
+                'notnull': data != mv,
+            },
             initial_workspace={c: data},
             mask=self.build_mask(self.ones_mask(shape=data.shape)),
         )
 
-        check_arrays(results['isnull'], (data == mv))
-        check_arrays(results['notnull'], (data != mv))
+    @parameter_space(compval=[0, 1, 999])
+    def test_eq(self, compval):
+
+        class C(Classifier):
+            dtype = int64_dtype
+            missing_value = -1
+            inputs = ()
+            window_length = 0
+
+        c = C()
+
+        # There's no significance to the values here other than that they
+        # contain a mix of the comparison value and other values.
+        data = np.array([[-1,  1,  0, 2],
+                         [3,   0,  1, 0],
+                         [-5,  0, -1, 0],
+                         [-3,  1,  2, 2]], dtype=int64_dtype)
+
+        self.check_terms(
+            terms={
+                'eq': c.eq(compval),
+            },
+            expected={
+                'eq': (data == compval),
+            },
+            initial_workspace={c: data},
+            mask=self.build_mask(self.ones_mask(shape=data.shape)),
+        )
+
+    @parameter_space(missing=[-1, 0, 1])
+    def test_disallow_comparison_to_missing_value(self, missing):
+        class C(Classifier):
+            dtype = int64_dtype
+            missing_value = missing
+            inputs = ()
+            window_length = 0
+
+        with self.assertRaises(ValueError) as e:
+            C().eq(missing)
+        errmsg = str(e.exception)
+        self.assertEqual(
+            errmsg,
+            "Comparison against self.missing_value ({v}) in C.eq().\n"
+            "Missing values have NaN semantics, so the requested comparison"
+            " would always produce False.\n"
+            "Use the isnull() method to check for missing values.".format(
+                v=missing,
+            ),
+        )
+
+    @parameter_space(compval=[0, 1, 999], missing=[-1, 0, 999])
+    def test_not_equal(self, compval, missing):
+
+        class C(Classifier):
+            dtype = int64_dtype
+            missing_value = missing
+            inputs = ()
+            window_length = 0
+
+        c = C()
+
+        # There's no significance to the values here other than that they
+        # contain a mix of the comparison value and other values.
+        data = np.array([[-1,  1,  0, 2],
+                         [3,   0,  1, 0],
+                         [-5,  0, -1, 0],
+                         [-3,  1,  2, 2]], dtype=int64_dtype)
+
+        self.check_terms(
+            terms={
+                'ne': c != compval,
+            },
+            expected={
+                'ne': (data != compval) & (data != C.missing_value),
+            },
+            initial_workspace={c: data},
+            mask=self.build_mask(self.ones_mask(shape=data.shape)),
+        )

zipline/pipeline/classifiers/classifier.py

@@ -1,13 +1,16 @@
 """
 classifier.py
 """
+from numbers import Number
+
 from numpy import where, isnan, nan, zeros
 
 from zipline.lib.quantiles import quantiles
 from zipline.pipeline.term import ComputableTerm
+from zipline.utils.input_validation import expect_types
 from zipline.utils.numpy_utils import int64_dtype
 
-from ..filters import NullFilter
+from ..filters import NullFilter, NumExprFilter
 from ..mixins import (
     CustomTermMixin,
     LatestMixin,
@@ -41,6 +44,48 @@ class Classifier(RestrictedDTypeMixin, ComputableTerm):
         """
         return ~self.isnull()
 
+    # We explicitly don't support classifier to classifier comparisons, since
+    # the numbers likely don't mean the same thing. This may be relaxed in the
+    # future, but for now we're starting conservatively.
+    @expect_types(other=Number)
+    def eq(self, other):
+        """
+        Construct a Filter returning True for asset/date pairs where the output
+        of ``self`` matches ``other.
+        """
+        # We treat this as an error because missing_values have NaN semantics,
+        # which means this would return an array of all False, which is almost
+        # certainly not what the user wants.
+        if other == self.missing_value:
+            raise ValueError(
+                "Comparison against self.missing_value ({value}) in"
+                " {typename}.eq().\n"
+                "Missing values have NaN semantics, so the "
+                "requested comparison would always produce False.\n"
+                "Use the isnull() method to check for missing values.".format(
+                    value=other,
+                    typename=(type(self).__name__),
+                )
+            )
+        return NumExprFilter.create(
+            "x_0 == {other}".format(other=int(other)),
+            binds=(self,),
+        )
+
+    @expect_types(other=Number)
+    def __ne__(self, other):
+        """
+        Construct a Filter returning True for asset/date pairs where the output
+        of ``self`` matches ``other.
+        """
+        return NumExprFilter.create(
+            "((x_0 != {other}) & (x_0 != {missing}))".format(
+                other=int(other),
+                missing=self.missing_value,
+            ),
+            binds=(self,),
+        )
+
 
 class Everything(Classifier):
     """

zipline/pipeline/expression.py

@@ -8,7 +8,7 @@ from numbers import Number
 import numexpr
 from numexpr.necompiler import getExprNames
 from numpy import (
-    empty,
+    full,
     inf,
 )
 
@@ -229,7 +229,7 @@ class NumericalExpression(ComputableTerm):
         """
         Compute our stored expression string with numexpr.
         """
-        out = empty(mask.shape, dtype=self.dtype)
+        out = full(mask.shape, self.missing_value, dtype=self.dtype)
         # This writes directly into our output buffer.
         numexpr.evaluate(
             self._expr,

zipline/pipeline/filters/filter.py

@@ -82,7 +82,7 @@ def binary_operator(op):
             )
         elif isinstance(other, int):  # Note that this is true for bool as well
             return NumExprFilter.create(
-                "x_0 {op} ({constant})".format(op=op, constant=int(other)),
+                "x_0 {op} {constant}".format(op=op, constant=int(other)),
                 binds=(self,),
             )
         raise BadBinaryOperator(op, self, other)