Merge "Ensure custom ops have consistent typing behavior, boolean support"

4 files changed, 69 insertions, 12 deletions

diff --git a/lib/sqlalchemy/sql/default_comparator.py b/lib/sqlalchemy/sql/default_comparator.py
index 4485c661b..a52bdbaed 100644
--- a/lib/sqlalchemy/sql/default_comparator.py
+++ b/lib/sqlalchemy/sql/default_comparator.py
@@ -81,6 +81,18 @@ def _boolean_compare(expr, op, obj, negate=None, reverse=False,
                                 negate=negate, modifiers=kwargs)
 
 
+def _custom_op_operate(expr, op, obj, reverse=False, result_type=None,
+                       **kw):
+    if result_type is None:
+        if op.return_type:
+            result_type = op.return_type
+        elif op.is_comparison:
+            result_type = type_api.BOOLEANTYPE
+
+    return _binary_operate(
+        expr, op, obj, reverse=reverse, result_type=result_type, **kw)
+
+
 def _binary_operate(expr, op, obj, reverse=False, result_type=None,
                     **kw):
     obj = _check_literal(expr, op, obj)
@@ -249,7 +261,7 @@ operator_lookup = {
     "div": (_binary_operate,),
     "mod": (_binary_operate,),
     "truediv": (_binary_operate,),
-    "custom_op": (_binary_operate,),
+    "custom_op": (_custom_op_operate,),
     "json_path_getitem_op": (_binary_operate, ),
     "json_getitem_op": (_binary_operate, ),
     "concat_op": (_binary_operate,),
diff --git a/lib/sqlalchemy/sql/operators.py b/lib/sqlalchemy/sql/operators.py
index f8731385b..a14afcb70 100644
--- a/lib/sqlalchemy/sql/operators.py
+++ b/lib/sqlalchemy/sql/operators.py
@@ -104,7 +104,9 @@ class Operators(object):
         """
         return self.operate(inv)
 
-    def op(self, opstring, precedence=0, is_comparison=False):
+    def op(
+            self, opstring, precedence=0, is_comparison=False,
+            return_type=None):
         """produce a generic operator function.
 
         e.g.::
@@ -145,6 +147,16 @@ class Operators(object):
          .. versionadded:: 0.9.2 - added the
             :paramref:`.Operators.op.is_comparison` flag.
 
+        :param return_type: a :class:`.TypeEngine` class or object that will
+          force the return type of an expression produced by this operator
+          to be of that type.   By default, operators that specify
+          :paramref:`.Operators.op.is_comparison` will resolve to
+          :class:`.Boolean`, and those that do not will be of the same
+          type as the left-hand operand.
+
+          .. versionadded:: 1.2.0b3 - added the
+             :paramref:`.Operators.op.return_type` argument.
+
         .. seealso::
 
             :ref:`types_operators`
@@ -152,12 +164,29 @@ class Operators(object):
             :ref:`relationship_custom_operator`
 
         """
-        operator = custom_op(opstring, precedence, is_comparison)
+        operator = custom_op(opstring, precedence, is_comparison, return_type)
 
         def against(other):
             return operator(self, other)
         return against
 
+    def bool_op(self, opstring, precedence=0):
+        """Return a custom boolean operator.
+
+        This method is shorthand for calling
+        :meth:`.Operators.op` and passing the
+        :paramref:`.Operators.op.is_comparison`
+        flag with True.
+
+        .. versionadded:: 1.2.0b3
+
+        .. seealso::
+
+            :meth:`.Operators.op`
+
+        """
+        return self.op(opstring, precedence=precedence, is_comparison=True)
+
     def operate(self, op, *other, **kwargs):
         r"""Operate on an argument.
 
@@ -197,9 +226,9 @@ class custom_op(object):
     """Represent a 'custom' operator.
 
     :class:`.custom_op` is normally instantiated when the
-    :meth:`.ColumnOperators.op` method is used to create a
-    custom operator callable.  The class can also be used directly
-    when programmatically constructing expressions.   E.g.
+    :meth:`.Operators.op` or :meth:`.Operators.bool_op` methods
+    are used to create a custom operator callable.  The class can also be
+    used directly when programmatically constructing expressions.   E.g.
     to represent the "factorial" operation::
 
         from sqlalchemy.sql import UnaryExpression
@@ -210,17 +239,28 @@ class custom_op(object):
                 modifier=operators.custom_op("!"),
                 type_=Numeric)
 
+
+    .. seealso::
+
+        :meth:`.Operators.op`
+
+        :meth:`.Operators.bool_op`
+
     """
     __name__ = 'custom_op'
 
     def __init__(
             self, opstring, precedence=0, is_comparison=False,
-            natural_self_precedent=False, eager_grouping=False):
+            return_type=None, natural_self_precedent=False,
+            eager_grouping=False):
         self.opstring = opstring
         self.precedence = precedence
         self.is_comparison = is_comparison
         self.natural_self_precedent = natural_self_precedent
         self.eager_grouping = eager_grouping
+        self.return_type = (
+            return_type._to_instance(return_type) if return_type else None
+        )
 
     def __eq__(self, other):
         return isinstance(other, custom_op) and \
diff --git a/lib/sqlalchemy/sql/sqltypes.py b/lib/sqlalchemy/sql/sqltypes.py
index d0dbc4881..5e357d39b 100644
--- a/lib/sqlalchemy/sql/sqltypes.py
+++ b/lib/sqlalchemy/sql/sqltypes.py
@@ -51,7 +51,7 @@ class _LookupExpressionAdapter(object):
             othertype = other_comparator.type._type_affinity
             lookup = self.type._expression_adaptations.get(
                 op, self._blank_dict).get(
-                othertype, NULLTYPE)
+                othertype, self.type)
             if lookup is othertype:
                 return (op, other_comparator.type)
             elif lookup is self.type._type_affinity:
@@ -2571,9 +2571,7 @@ class NullType(TypeEngine):
     class Comparator(TypeEngine.Comparator):
 
         def _adapt_expression(self, op, other_comparator):
-            if operators.is_comparison(op):
-                return op, BOOLEANTYPE
-            elif isinstance(other_comparator, NullType.Comparator) or \
+            if isinstance(other_comparator, NullType.Comparator) or \
                     not operators.is_commutative(op):
                 return op, self.expr.type
             else:
diff --git a/lib/sqlalchemy/sql/type_api.py b/lib/sqlalchemy/sql/type_api.py
index 4b561a705..69dd80938 100644
--- a/lib/sqlalchemy/sql/type_api.py
+++ b/lib/sqlalchemy/sql/type_api.py
@@ -93,6 +93,7 @@ class TypeEngine(Visitable):
             boolean comparison or special SQL keywords like MATCH or BETWEEN.
 
             """
+
             return op, self.type
 
         def __reduce__(self):
@@ -353,6 +354,10 @@ class TypeEngine(Visitable):
         return self.__class__.bind_expression.__code__ \
             is not TypeEngine.bind_expression.__code__
 
+    @staticmethod
+    def _to_instance(cls_or_self):
+        return to_instance(cls_or_self)
+
     def compare_values(self, x, y):
         """Compare two values for equality."""
 
@@ -634,7 +639,9 @@ class UserDefinedType(util.with_metaclass(VisitableCheckKWArg, TypeEngine)):
                 )
                 return self.type.adapt_operator(op), self.type
             else:
-                return op, self.type
+                return super(
+                    UserDefinedType.Comparator, self
+                )._adapt_expression(op, other_comparator)
 
     comparator_factory = Comparator