Implement @total_ordering decorator for extension types (GH-3626)

Implements https://github.com/cython/cython/issues/2090

Cython/Compiler/ModuleNode.py

@@ -72,6 +72,41 @@ def generate_c_code_config(env, options):
         emit_code_comments=env.directives['emit_code_comments'],
         c_line_in_traceback=options.c_line_in_traceback)
 
+# The code required to generate one comparison from another.
+# The keys are (from, to).
+# The comparison operator always goes first, with equality possibly second.
+# The first value specifies if the comparison is inverted. The second is the
+# logic op to use, and the third is if the equality is inverted or not.
+TOTAL_ORDERING = {
+    # a > b from (not a < b) and (a != b)
+    ('__lt__', '__gt__'): (True, '&&', True),
+    # a <= b from (a < b) or (a == b)
+    ('__lt__', '__le__'): (False, '||', False),
+    # a >= b from (not a < b).
+    ('__lt__', '__ge__'): (True, '', None),
+
+    # a >= b from (not a <= b) or (a == b)
+    ('__le__', '__ge__'): (True, '||', False),
+    # a < b, from (a <= b) and (a != b)
+    ('__le__', '__lt__'): (False, '&&', True),
+    # a > b from (not a <= b)
+    ('__le__', '__gt__'): (True, '', None),
+
+    # a < b from (not a > b) and (a != b)
+    ('__gt__', '__lt__'): (True, '&&', True),
+    # a >= b from (a > b) or (a == b)
+    ('__gt__', '__ge__'): (False, '||', False),
+    # a <= b from (not a > b)
+    ('__gt__', '__le__'): (True, '', None),
+
+    # Return a <= b from (not a >= b) or (a == b)
+    ('__ge__', '__le__'): (True, '||', False),
+    # a > b from (a >= b) and (a != b)
+    ('__ge__', '__gt__'): (False, '&&', True),
+    # a < b from (not a >= b)
+    ('__ge__', '__lt__'): (True, '', None),
+}
+
 
 class ModuleNode(Nodes.Node, Nodes.BlockNode):
     #  doc       string or None
@@ -1356,6 +1391,7 @@ class ModuleNode(Nodes.Node, Nodes.BlockNode):
                     self.generate_exttype_vtable(scope, code)
                     self.generate_new_function(scope, code, entry)
                     self.generate_dealloc_function(scope, code)
+
                     if scope.needs_gc():
                         self.generate_traverse_function(scope, code, entry)
                         if scope.needs_tp_clear():
@@ -1383,11 +1419,18 @@ class ModuleNode(Nodes.Node, Nodes.BlockNode):
                         self.generate_descr_set_function(scope, code)
                     if not scope.is_closure_class_scope and scope.defines_any(["__dict__"]):
                         self.generate_dict_getter_function(scope, code)
+
                     if scope.defines_any_special(TypeSlots.richcmp_special_methods):
                         self.generate_richcmp_function(scope, code)
+                    elif scope.directives.get('total_ordering'):
+                        # Warn if this is used when it can't have any effect.
+                        warning(scope.parent_type.pos,
+                                "total_ordering directive used, but no comparison and equality methods defined")
+
                     for slot in TypeSlots.PyNumberMethods:
                         if slot.is_binop and scope.defines_any_special(slot.user_methods):
                             self.generate_binop_function(scope, slot, code, entry.pos)
+
                     self.generate_property_accessors(scope, code)
                     self.generate_method_table(scope, code)
                     self.generate_getset_table(scope, code)
@@ -2035,37 +2078,112 @@ class ModuleNode(Nodes.Node, Nodes.BlockNode):
             # need to call up into base classes as we may not know all implemented comparison methods
             extern_parent = cls if cls.typeptr_cname else scope.parent_type.base_type
 
-        eq_entry = None
-        has_ne = False
+        total_ordering = scope.directives.get('total_ordering', False)
+
+        comp_entry = {}
+
         for cmp_method in TypeSlots.richcmp_special_methods:
             for class_scope in class_scopes:
                 entry = class_scope.lookup_here(cmp_method)
                 if entry is not None:
+                    comp_entry[cmp_method] = entry
                     break
+
+        if total_ordering:
+            # Check this is valid - we must have at least 1 operation defined.
+            comp_names = [from_name for from_name, to_name in TOTAL_ORDERING if from_name in comp_entry]
+            if not comp_names:
+                if '__eq__' not in comp_entry and '__ne__'  not in comp_entry:
+                    warning(scope.parent_type.pos,
+                            "total_ordering directive used, but no comparison and equality methods defined")
+                else:
+                    warning(scope.parent_type.pos,
+                          "total_ordering directive used, but no comparison methods defined")
+                total_ordering = False
             else:
-                continue
+                if '__eq__' not in comp_entry and '__ne__' not in comp_entry:
+                    warning(scope.parent_type.pos, "total_ordering directive used, but no equality method defined")
+                    total_ordering = False
+
+                # Same priority as functools, prefers
+                # __lt__ to __le__ to __gt__ to __ge__
+                ordering_source = max(comp_names)
 
+        for cmp_method in TypeSlots.richcmp_special_methods:
             cmp_type = cmp_method.strip('_').upper()  # e.g. "__eq__" -> EQ
+            entry = comp_entry.get(cmp_method)
+            if entry is None and (not total_ordering or cmp_type in ('NE', 'EQ')):
+                # No definition, fall back to superclasses.
+                # eq/ne methods shouldn't use the total_ordering code.
+                continue
+
             code.putln("case Py_%s: {" % cmp_type)
-            if cmp_method == '__eq__':
-                eq_entry = entry
-                # Python itself does not do this optimisation, it seems...
-                #code.putln("if (o1 == o2) return __Pyx_NewRef(Py_True);")
-            elif cmp_method == '__ne__':
-                has_ne = True
-                # Python itself does not do this optimisation, it seems...
-                #code.putln("if (o1 == o2) return __Pyx_NewRef(Py_False);")
-            code.putln("return %s(o1, o2);" % entry.func_cname)
-            code.putln("}")
+            if entry is None:
+                assert total_ordering
+                # We need to generate this from the other methods.
+                invert_comp, comp_op, invert_equals = TOTAL_ORDERING[ordering_source, cmp_method]
+
+                # First we always do the comparison.
+                code.putln("PyObject *ret;")
+                code.putln("ret = %s(o1, o2);" % comp_entry[ordering_source].func_cname)
+                code.putln("if (likely(ret && ret != Py_NotImplemented)) {")
+                code.putln("int order_res = __Pyx_PyObject_IsTrue(ret);")
+                code.putln("Py_DECREF(ret);")
+                code.putln("if (unlikely(order_res < 0)) return NULL;")
+                # We may need to check equality too. For some combos it's never required.
+                if invert_equals is not None:
+                    # Implement the and/or check with an if.
+                    if comp_op == '&&':
+                        code.putln("if (%s order_res) {" % ('!!' if invert_comp else '!'))
+                        code.putln("ret = Py_False;")
+                        code.putln("} else {")
+                    elif comp_op == '||':
+                        code.putln("if (%s order_res) {" % ('!' if invert_comp else ''))
+                        code.putln("ret = Py_True;")
+                        code.putln("} else {")
+                    else:
+                        raise AssertionError('Unknown op %s' % (comp_op, ))
+                    if '__eq__' in comp_entry:
+                        eq_func = '__eq__'
+                    else:
+                        # Fall back to NE, which is defined here.
+                        eq_func = '__ne__'
+                        invert_equals = not invert_equals
+
+                    code.putln("ret = %s(o1, o2);" % comp_entry[eq_func].func_cname)
+                    code.putln("if (likely(ret && ret != Py_NotImplemented)) {")
+                    code.putln("int eq_res = __Pyx_PyObject_IsTrue(ret);")
+                    code.putln("Py_DECREF(ret);")
+                    code.putln("if (unlikely(eq_res < 0)) return NULL;")
+                    if invert_equals:
+                        code.putln("ret = eq_res ? Py_False : Py_True;")
+                    else:
+                        code.putln("ret = eq_res ? Py_True : Py_False;")
+                    code.putln("Py_INCREF(ret);")
+                    code.putln("}")  # equals success
+                    code.putln("}")  # Needs to try equals
+                else:
+                    # Convert direct to a boolean.
+                    if invert_comp:
+                        code.putln("ret = order_res ? Py_False : Py_True;")
+                    else:
+                        code.putln("ret = order_res ? Py_True : Py_False;")
+                    code.putln("Py_INCREF(ret);")
+                code.putln("}")  # comp_op
+                code.putln("return ret;")
+            else:
+                code.putln("return %s(o1, o2);" % entry.func_cname)
+            code.putln("}")  # Case
 
-        if eq_entry and not has_ne and not extern_parent:
+        if '__eq__' in comp_entry and '__ne__' not in comp_entry and not extern_parent:
             code.putln("case Py_NE: {")
             code.putln("PyObject *ret;")
             # Python itself does not do this optimisation, it seems...
             #code.putln("if (o1 == o2) return __Pyx_NewRef(Py_False);")
-            code.putln("ret = %s(o1, o2);" % eq_entry.func_cname)
+            code.putln("ret = %s(o1, o2);" % comp_entry['__eq__'].func_cname)
             code.putln("if (likely(ret && ret != Py_NotImplemented)) {")
-            code.putln("int b = __Pyx_PyObject_IsTrue(ret); Py_DECREF(ret);")
+            code.putln("int b = __Pyx_PyObject_IsTrue(ret);")
+            code.putln("Py_DECREF(ret);")
             code.putln("if (unlikely(b < 0)) return NULL;")
             code.putln("ret = (b) ? Py_False : Py_True;")
             code.putln("Py_INCREF(ret);")

Cython/Compiler/Options.py

@@ -326,6 +326,7 @@ directive_types = {
     'c_string_type': one_of('bytes', 'bytearray', 'str', 'unicode'),
     'c_string_encoding': normalise_encoding_name,
     'trashcan': bool,
+    'total_ordering': bool,
 }
 
 for key, val in _directive_defaults.items():
@@ -369,6 +370,7 @@ directive_scopes = {  # defaults to available everywhere
     'fast_gil': ('module',),
     'iterable_coroutine': ('module', 'function'),
     'trashcan' : ('cclass',),
+    'total_ordering': ('cclass', ),
 }
 
 

Cython/Shadow.py

@@ -121,7 +121,7 @@ optimization = _Optimization()
 overflowcheck.fold = optimization.use_switch = \
     optimization.unpack_method_calls = lambda arg: _EmptyDecoratorAndManager()
 
-final = internal = type_version_tag = no_gc_clear = no_gc = _empty_decorator
+final = internal = type_version_tag = no_gc_clear = no_gc = total_ordering = _empty_decorator
 
 binding = lambda _: _empty_decorator
 

tests/errors/e_exttype_total_ordering.pyx

+# mode: error
+# tag: total_ordering, warnings
+
+cimport cython
+
+
+# Test all combinations with not enough methods.
+
+@cython.total_ordering
+cdef class ExtNoFuncs:
+    pass
+
+@cython.total_ordering
+cdef class ExtGe:
+    def __ge__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtLe:
+    def __le__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtLeGe:
+    def __le__(self, other):
+        return False
+
+    def __ge__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtGt:
+    def __gt__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtGtGe:
+    def __gt__(self, other):
+        return False
+
+    def __ge__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtGtLe:
+    def __gt__(self, other):
+        return False
+
+    def __le__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtGtLeGe:
+    def __gt__(self, other):
+        return False
+
+    def __le__(self, other):
+        return False
+
+    def __ge__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtLt:
+    def __lt__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtLtGe:
+    def __lt__(self, other):
+        return False
+
+    def __ge__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtLtLe:
+    def __lt__(self, other):
+        return False
+
+    def __le__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtLtLeGe:
+    def __lt__(self, other):
+        return False
+
+    def __le__(self, other):
+        return False
+
+    def __ge__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtLtGt:
+    def __lt__(self, other):
+        return False
+
+    def __gt__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtLtGtGe:
+    def __lt__(self, other):
+        return False
+
+    def __gt__(self, other):
+        return False
+
+    def __ge__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtLtGtLe:
+    def __lt__(self, other):
+        return False
+
+    def __gt__(self, other):
+        return False
+
+    def __le__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtLtGtLeGe:
+    def __lt__(self, other):
+        return False
+
+    def __gt__(self, other):
+        return False
+
+    def __le__(self, other):
+        return False
+
+    def __ge__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtNe:
+    def __ne__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtEq:
+    def __eq__(self, other):
+        return False
+
+@cython.total_ordering
+cdef class ExtEqNe:
+    def __eq__(self, other):
+        return False
+
+    def __ne__(self, other):
+        return False
+
+
+_WARNINGS = """
+10:5: total_ordering directive used, but no comparison and equality methods defined
+14:5: total_ordering directive used, but no equality method defined
+19:5: total_ordering directive used, but no equality method defined
+24:5: total_ordering directive used, but no equality method defined
+32:5: total_ordering directive used, but no equality method defined
+37:5: total_ordering directive used, but no equality method defined
+45:5: total_ordering directive used, but no equality method defined
+53:5: total_ordering directive used, but no equality method defined
+64:5: total_ordering directive used, but no equality method defined
+69:5: total_ordering directive used, but no equality method defined
+77:5: total_ordering directive used, but no equality method defined
+85:5: total_ordering directive used, but no equality method defined
+96:5: total_ordering directive used, but no equality method defined
+104:5: total_ordering directive used, but no equality method defined
+115:5: total_ordering directive used, but no equality method defined
+126:5: total_ordering directive used, but no equality method defined
+140:5: total_ordering directive used, but no comparison methods defined
+145:5: total_ordering directive used, but no comparison methods defined
+150:5: total_ordering directive used, but no comparison methods defined
+"""