GR-10939: sets -> add support for comparison with objects implementing the richcmp protocol

- added relevant unittest

Author: cosminbasca

graalpython/com.oracle.graal.python.test/src/tests/test_set.py

@@ -215,3 +215,40 @@ def test_intersection():
 def test_same_id():
     empty_ids = set([id(frozenset()) for i in range(100)])
     assert len(empty_ids) == 1
+
+
+def test_rich_compare():
+    class TestRichSetCompare:
+        def __gt__(self, some_set):
+            self.gt_called = True
+            return False
+        def __lt__(self, some_set):
+            self.lt_called = True
+            return False
+        def __ge__(self, some_set):
+            self.ge_called = True
+            return False
+        def __le__(self, some_set):
+            self.le_called = True
+            return False
+
+    # This first tries the builtin rich set comparison, which doesn't know
+    # how to handle the custom object. Upon returning NotImplemented, the
+    # corresponding comparison on the right object is invoked.
+    myset = {1, 2, 3}
+
+    myobj = TestRichSetCompare()
+    myset < myobj
+    assert myobj.gt_called
+
+    myobj = TestRichSetCompare()
+    myset > myobj
+    assert myobj.lt_called
+
+    myobj = TestRichSetCompare()
+    myset <= myobj
+    assert myobj.ge_called
+
+    myobj = TestRichSetCompare()
+    myset >= myobj
+    assert myobj.le_called

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/set/FrozenSetBuiltins.java

@@ -53,13 +53,15 @@
 import com.oracle.graal.python.builtins.objects.common.PHashingCollection;
 import com.oracle.graal.python.builtins.objects.set.FrozenSetBuiltinsFactory.BinaryUnionNodeGen;
 import com.oracle.graal.python.nodes.PBaseNode;
+import com.oracle.graal.python.nodes.call.special.LookupAndCallBinaryNode;
 import com.oracle.graal.python.nodes.control.GetIteratorNode;
 import com.oracle.graal.python.nodes.control.GetNextNode;
 import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
 import com.oracle.graal.python.nodes.function.PythonBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonBinaryBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonUnaryBuiltinNode;
 import com.oracle.graal.python.runtime.exception.PException;
+import com.oracle.graal.python.runtime.exception.PythonErrorType;
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
@@ -366,11 +368,29 @@ boolean isSuperSet(PBaseSet self, String other,
     @Builtin(name = __LE__, fixedNumOfArguments = 2)
     @GenerateNodeFactory
     abstract static class LessEqualNode extends IsSubsetNode {
+        @Specialization
+        Object isLessEqual(PBaseSet self, Object other,
+                        @Cached("create(__GE__)") LookupAndCallBinaryNode lookupAndCallBinaryNode) {
+            Object result = lookupAndCallBinaryNode.executeObject(other, self);
+            if (result != PNone.NO_VALUE) {
+                return result;
+            }
+            throw raise(PythonErrorType.TypeError, "unorderable types: %p <= %p", self, other);
+        }
     }
 
     @Builtin(name = __GE__, fixedNumOfArguments = 2)
     @GenerateNodeFactory
     abstract static class GreaterEqualNode extends IsSupersetNode {
+        @Specialization
+        Object isGreaterEqual(PBaseSet self, Object other,
+                        @Cached("create(__LE__)") LookupAndCallBinaryNode lookupAndCallBinaryNode) {
+            Object result = lookupAndCallBinaryNode.executeObject(other, self);
+            if (result != PNone.NO_VALUE) {
+                return result;
+            }
+            throw raise(PythonErrorType.TypeError, "unorderable types: %p >= %p", self, other);
+        }
     }
 
     @Builtin(name = __LT__, fixedNumOfArguments = 2)
@@ -403,6 +423,16 @@ boolean isLessThan(PBaseSet self, String other,
             }
             return (Boolean) getLessEqualNode().execute(self, other);
         }
+
+        @Specialization
+        Object isLessThan(PBaseSet self, Object other,
+                        @Cached("create(__GT__)") LookupAndCallBinaryNode lookupAndCallBinaryNode) {
+            Object result = lookupAndCallBinaryNode.executeObject(other, self);
+            if (result != PNone.NO_VALUE) {
+                return result;
+            }
+            throw raise(PythonErrorType.TypeError, "unorderable types: %p < %p", self, other);
+        }
     }
 
     @Builtin(name = __GT__, fixedNumOfArguments = 2)
@@ -435,5 +465,15 @@ boolean isGreaterThan(PBaseSet self, String other,
             }
             return (Boolean) getGreaterEqualNode().execute(self, other);
         }
+
+        @Specialization
+        Object isLessThan(PBaseSet self, Object other,
+                        @Cached("create(__LT__)") LookupAndCallBinaryNode lookupAndCallBinaryNode) {
+            Object result = lookupAndCallBinaryNode.executeObject(other, self);
+            if (result != PNone.NO_VALUE) {
+                return result;
+            }
+            throw raise(PythonErrorType.TypeError, "unorderable types: %p > %p", self, other);
+        }
     }
 }