From acfdf9ceaa2a67924457dffd3e6e4f124bbafa3e Mon Sep 17 00:00:00 2001
From: jfh <jfh@imijmi.eu>
Date: Thu, 7 Oct 2021 12:23:12 +0300
Subject: [PATCH] Move try_complex into PyObjectRef

---
 vm/src/builtins/complex.rs | 65 +++++++++++++++++++-------------------
 vm/src/function/numlike.rs |  7 ++--
 2 files changed, 34 insertions(+), 38 deletions(-)

diff --git a/vm/src/builtins/complex.rs b/vm/src/builtins/complex.rs
index 223f29064..671979171 100644
--- a/vm/src/builtins/complex.rs
+++ b/vm/src/builtins/complex.rs
@@ -39,6 +39,36 @@ impl From<Complex64> for PyComplex {
     }
 }
 
+impl PyObjectRef {
+    /// Tries converting a python object into a complex, returns an option of whether the complex
+    /// and whether the  object was a complex originally or coereced into one
+    pub fn try_complex(&self, vm: &VirtualMachine) -> PyResult<Option<(Complex64, bool)>> {
+        if let Some(complex) = self.payload_if_exact::<PyComplex>(vm) {
+            return Ok(Some((complex.value, true)));
+        }
+        if let Some(method) = vm.get_method(self.clone(), "__complex__") {
+            let result = vm.invoke(&method?, ())?;
+            // TODO: returning strict subclasses of complex in __complex__ is deprecated
+            return match result.payload::<PyComplex>() {
+                Some(complex_obj) => Ok(Some((complex_obj.value, true))),
+                None => Err(vm.new_type_error(format!(
+                    "__complex__ returned non-complex (type '{}')",
+                    result.class().name()
+                ))),
+            };
+        }
+        // `complex` does not have a `__complex__` by default, so subclasses might not either,
+        // use the actual stored value in this case
+        if let Some(complex) = self.payload_if_subclass::<PyComplex>(vm) {
+            return Ok(Some((complex.value, true)));
+        }
+        if let Some(float) = self.try_to_f64(vm)? {
+            return Ok(Some((Complex64::new(float, 0.0), false)));
+        }
+        Ok(None)
+    }
+}
+
 pub fn init(context: &PyContext) {
     PyComplex::extend_class(context, &context.types.complex_type);
 }
@@ -99,7 +129,7 @@ impl SlotConstructor for PyComplex {
                     val
                 };
 
-                if let Some(c) = try_complex(&val, vm)? {
+                if let Some(c) = val.try_complex(vm)? {
                     c
                 } else if let Some(s) = val.payload_if_subclass::<PyStr>(vm) {
                     if args.imag.is_present() {
@@ -125,7 +155,7 @@ impl SlotConstructor for PyComplex {
             // if an  imaginary argument is not passed in
             OptionalArg::Missing => (Complex64::new(real.im, 0.0), false),
             OptionalArg::Present(obj) => {
-                if let Some(c) = try_complex(&obj, vm)? {
+                if let Some(c) = obj.try_complex(vm)? {
                     c
                 } else if obj.class().issubclass(&vm.ctx.types.str_type) {
                     return Err(
@@ -420,34 +450,3 @@ fn parse_str(s: &str) -> Option<Complex64> {
     };
     Some(value)
 }
-
-/// Tries converting a python object into a complex, returns an option of whether the complex
-/// and whether the  object was a complex originally or coereced into one
-pub(crate) fn try_complex(
-    obj: &PyObjectRef,
-    vm: &VirtualMachine,
-) -> PyResult<Option<(Complex64, bool)>> {
-    if let Some(complex) = obj.payload_if_exact::<PyComplex>(vm) {
-        return Ok(Some((complex.value, true)));
-    }
-    if let Some(method) = vm.get_method(obj.clone(), "__complex__") {
-        let result = vm.invoke(&method?, ())?;
-        // TODO: returning strict subclasses of complex in __complex__ is deprecated
-        return match result.payload::<PyComplex>() {
-            Some(complex_obj) => Ok(Some((complex_obj.value, true))),
-            None => Err(vm.new_type_error(format!(
-                "__complex__ returned non-complex (type '{}')",
-                result.class().name()
-            ))),
-        };
-    }
-    // `complex` does not have a `__complex__` by default, so subclasses might not either,
-    // use the actual stored value in this case
-    if let Some(complex) = obj.payload_if_subclass::<PyComplex>(vm) {
-        return Ok(Some((complex.value, true)));
-    }
-    if let Some(float) = obj.try_to_f64(vm)? {
-        return Ok(Some((Complex64::new(float, 0.0), false)));
-    }
-    Ok(None)
-}
diff --git a/vm/src/function/numlike.rs b/vm/src/function/numlike.rs
index aa18f6079..838c7ec40 100644
--- a/vm/src/function/numlike.rs
+++ b/vm/src/function/numlike.rs
@@ -1,7 +1,4 @@
-use crate::{
-    builtins::complex::try_complex, PyObjectRef, PyResult, TryFromObject, TypeProtocol,
-    VirtualMachine,
-};
+use crate::{PyObjectRef, PyResult, TryFromObject, TypeProtocol, VirtualMachine};
 use num_complex::Complex64;
 
 #[derive(Debug, Copy, Clone, PartialEq)]
@@ -19,7 +16,7 @@ impl ArgComplexLike {
 impl TryFromObject for ArgComplexLike {
     fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
         // We do not care if it was already a complex.
-        let (value, _) = try_complex(&obj, vm)?.ok_or_else(|| {
+        let (value, _) = obj.try_complex(vm)?.ok_or_else(|| {
             vm.new_type_error(format!("must be real number, not {}", obj.class().name()))
         })?;
         Ok(ArgComplexLike { value })