Close #10042: functools.total_ordering now handles NotImplemented

(Patch by Katie Miller)

1 files changed, 101 insertions, 7 deletions

diff --git a/Lib/test/test_functools.py b/Lib/test/test_functools.py
index ab76efb..cb493bf 100644
--- a/Lib/test/test_functools.py
+++ b/Lib/test/test_functools.py
@@ -584,6 +584,7 @@ class TestTotalOrdering(unittest.TestCase):
         self.assertTrue(A(2) >= A(1))
         self.assertTrue(A(2) <= A(2))
         self.assertTrue(A(2) >= A(2))
+        self.assertFalse(A(1) > A(2))
 
     def test_total_ordering_le(self):
         @functools.total_ordering
@@ -600,6 +601,7 @@ class TestTotalOrdering(unittest.TestCase):
         self.assertTrue(A(2) >= A(1))
         self.assertTrue(A(2) <= A(2))
         self.assertTrue(A(2) >= A(2))
+        self.assertFalse(A(1) >= A(2))
 
     def test_total_ordering_gt(self):
         @functools.total_ordering
@@ -616,6 +618,7 @@ class TestTotalOrdering(unittest.TestCase):
         self.assertTrue(A(2) >= A(1))
         self.assertTrue(A(2) <= A(2))
         self.assertTrue(A(2) >= A(2))
+        self.assertFalse(A(2) < A(1))
 
     def test_total_ordering_ge(self):
         @functools.total_ordering
@@ -632,6 +635,7 @@ class TestTotalOrdering(unittest.TestCase):
         self.assertTrue(A(2) >= A(1))
         self.assertTrue(A(2) <= A(2))
         self.assertTrue(A(2) >= A(2))
+        self.assertFalse(A(2) <= A(1))
 
     def test_total_ordering_no_overwrite(self):
         # new methods should not overwrite existing
@@ -651,22 +655,112 @@ class TestTotalOrdering(unittest.TestCase):
             class A:
                 pass
 
-    def test_bug_10042(self):
+    def test_type_error_when_not_implemented(self):
+        # bug 10042; ensure stack overflow does not occur
+        # when decorated types return NotImplemented
         @functools.total_ordering
-        class TestTO:
+        class ImplementsLessThan:
             def __init__(self, value):
                 self.value = value
             def __eq__(self, other):
-                if isinstance(other, TestTO):
+                if isinstance(other, ImplementsLessThan):
                     return self.value == other.value
                 return False
             def __lt__(self, other):
-                if isinstance(other, TestTO):
+                if isinstance(other, ImplementsLessThan):
                     return self.value < other.value
-                raise TypeError
-        with self.assertRaises(TypeError):
-            TestTO(8) <= ()
+                return NotImplemented
+
+        @functools.total_ordering
+        class ImplementsGreaterThan:
+            def __init__(self, value):
+                self.value = value
+            def __eq__(self, other):
+                if isinstance(other, ImplementsGreaterThan):
+                    return self.value == other.value
+                return False
+            def __gt__(self, other):
+                if isinstance(other, ImplementsGreaterThan):
+                    return self.value > other.value
+                return NotImplemented
+
+        @functools.total_ordering
+        class ImplementsLessThanEqualTo:
+            def __init__(self, value):
+                self.value = value
+            def __eq__(self, other):
+                if isinstance(other, ImplementsLessThanEqualTo):
+                    return self.value == other.value
+                return False
+            def __le__(self, other):
+                if isinstance(other, ImplementsLessThanEqualTo):
+                    return self.value <= other.value
+                return NotImplemented
+
+        @functools.total_ordering
+        class ImplementsGreaterThanEqualTo:
+            def __init__(self, value):
+                self.value = value
+            def __eq__(self, other):
+                if isinstance(other, ImplementsGreaterThanEqualTo):
+                    return self.value == other.value
+                return False
+            def __ge__(self, other):
+                if isinstance(other, ImplementsGreaterThanEqualTo):
+                    return self.value >= other.value
+                return NotImplemented
+
+        @functools.total_ordering
+        class ComparatorNotImplemented:
+            def __init__(self, value):
+                self.value = value
+            def __eq__(self, other):
+                if isinstance(other, ComparatorNotImplemented):
+                    return self.value == other.value
+                return False
+            def __lt__(self, other):
+                return NotImplemented
+
+        with self.subTest("LT < 1"), self.assertRaises(TypeError):
+            ImplementsLessThan(-1) < 1
+
+        with self.subTest("LT < LE"), self.assertRaises(TypeError):
+            ImplementsLessThan(0) < ImplementsLessThanEqualTo(0)
+
+        with self.subTest("LT < GT"), self.assertRaises(TypeError):
+            ImplementsLessThan(1) < ImplementsGreaterThan(1)
+
+        with self.subTest("LE <= LT"), self.assertRaises(TypeError):
+            ImplementsLessThanEqualTo(2) <= ImplementsLessThan(2)
+
+        with self.subTest("LE <= GE"), self.assertRaises(TypeError):
+            ImplementsLessThanEqualTo(3) <= ImplementsGreaterThanEqualTo(3)
+
+        with self.subTest("GT > GE"), self.assertRaises(TypeError):
+            ImplementsGreaterThan(4) > ImplementsGreaterThanEqualTo(4)
+
+        with self.subTest("GT > LT"), self.assertRaises(TypeError):
+            ImplementsGreaterThan(5) > ImplementsLessThan(5)
+
+        with self.subTest("GE >= GT"), self.assertRaises(TypeError):
+            ImplementsGreaterThanEqualTo(6) >= ImplementsGreaterThan(6)
+
+        with self.subTest("GE >= LE"), self.assertRaises(TypeError):
+            ImplementsGreaterThanEqualTo(7) >= ImplementsLessThanEqualTo(7)
+
+        with self.subTest("GE when equal"):
+            a = ComparatorNotImplemented(8)
+            b = ComparatorNotImplemented(8)
+            self.assertEqual(a, b)
+            with self.assertRaises(TypeError):
+                a >= b
 
+        with self.subTest("LE when equal"):
+            a = ComparatorNotImplemented(9)
+            b = ComparatorNotImplemented(9)
+            self.assertEqual(a, b)
+            with self.assertRaises(TypeError):
+                a <= b
 
 class TestLRU(unittest.TestCase):