From 9f57c58dad917136f27cc0250971b7e942fc54d0 Mon Sep 17 00:00:00 2001
From: jfh <jfh@imijmi.eu>
Date: Wed, 9 Jun 2021 02:06:48 +0300
Subject: [PATCH] Handle NotImplemented in length_hint.

---
 Lib/test/test_operator.py |  5 -----
 vm/src/iterator.rs        |  9 +++++++--
 vm/src/stdlib/operator.rs | 25 ++++++++++++++-----------
 3 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/Lib/test/test_operator.py b/Lib/test/test_operator.py
index e260433c7..f46d94a22 100644
--- a/Lib/test/test_operator.py
+++ b/Lib/test/test_operator.py
@@ -516,11 +516,6 @@ class PyOperatorTestCase(OperatorTestCase, unittest.TestCase):
 class COperatorTestCase(OperatorTestCase, unittest.TestCase):
     module = c_operator
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
-    def test_length_hint(self):
-        super().test_length_hint()
-
 
 class OperatorPickleTestCase:
     def copy(self, obj, proto):
diff --git a/vm/src/iterator.rs b/vm/src/iterator.rs
index e344b9862..f8708fa02 100644
--- a/vm/src/iterator.rs
+++ b/vm/src/iterator.rs
@@ -6,7 +6,7 @@ use crate::builtins::int::{self, PyInt};
 use crate::builtins::iter::PySequenceIterator;
 use crate::exceptions::PyBaseExceptionRef;
 use crate::vm::VirtualMachine;
-use crate::{PyObjectRef, PyResult, PyValue, TryFromObject, TypeProtocol};
+use crate::{IdProtocol, PyObjectRef, PyResult, PyValue, TryFromObject, TypeProtocol};
 use num_traits::Signed;
 
 /*
@@ -122,7 +122,12 @@ pub fn length_hint(vm: &VirtualMachine, iter: PyObjectRef) -> PyResult<Option<us
         None => return Ok(None),
     };
     let result = match vm.invoke(&hint, ()) {
-        Ok(res) => res,
+        Ok(res) => {
+            if res.is(&vm.ctx.not_implemented) {
+                return Ok(None);
+            }
+            res
+        }
         Err(e) => {
             return if e.isinstance(&vm.ctx.exceptions.type_error) {
                 Ok(None)
diff --git a/vm/src/stdlib/operator.rs b/vm/src/stdlib/operator.rs
index 1d96423cd..ec33daf1f 100644
--- a/vm/src/stdlib/operator.rs
+++ b/vm/src/stdlib/operator.rs
@@ -1,3 +1,4 @@
+use crate::builtins::int;
 use crate::builtins::int::PyIntRef;
 use crate::builtins::pystr::PyStrRef;
 use crate::byteslike::PyBytesLike;
@@ -7,19 +8,21 @@ use crate::iterator;
 use crate::utils::Either;
 use crate::VirtualMachine;
 use crate::{PyObjectRef, PyResult, TypeProtocol};
+use int::PyInt;
 
 fn _operator_length_hint(obj: PyObjectRef, default: OptionalArg, vm: &VirtualMachine) -> PyResult {
-    let default = default.unwrap_or_else(|| vm.ctx.new_int(0));
-    if !default.isinstance(&vm.ctx.types.int_type) {
-        return Err(vm.new_type_error(format!(
-            "'{}' type cannot be interpreted as an integer",
-            default.class().name
-        )));
-    }
-    let hint = iterator::length_hint(vm, obj)?
-        .map(|i| vm.ctx.new_int(i))
-        .unwrap_or(default);
-    Ok(hint)
+    let default: usize = default
+        .map(|v| {
+            if !v.isinstance(&vm.ctx.types.int_type) {
+                return Err(vm.new_type_error(format!(
+                    "'{}' type cannot be interpreted as an integer",
+                    v.class().name
+                )));
+            }
+            int::try_to_primitive(v.payload::<PyInt>().unwrap().as_bigint(), vm)
+        })
+        .unwrap_or(Ok(0))?;
+    iterator::length_hint(vm, obj).map(|v| vm.ctx.new_int(v.unwrap_or(default)))
 }
 
 fn _operator_compare_digest(