Implement PyComplex's __format__ function (#5900)

* Add num-complex in rustpython-common

* Implement PyComplex's __format__ function

* Remove @unittest.expectedFailure

* Fix PyComplex's __format__ function

* Remove @unittest.expectedFailure

* Add extra tests

* Rename to ZeroPadding and AlignmentFlag

Cargo.lock

@@ -2300,6 +2300,7 @@ dependencies = [
  "malachite-bigint",
  "malachite-q",
  "memchr",
+ "num-complex",
  "num-traits",
  "once_cell",
  "parking_lot",

Lib/test/test_complex.py

@@ -719,8 +719,6 @@ class ComplexTest(unittest.TestCase):
                 self.assertFloatsAreIdentical(0.0 + z.imag,
                                               0.0 + roundtrip.imag)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_format(self):
         # empty format string is same as str()
         self.assertEqual(format(1+3j, ''), str(1+3j))

Lib/test/test_format.py

@@ -397,8 +397,6 @@ class FormatTest(unittest.TestCase):
         testformat("a%sb", ('c\0d',), 'ac\0db')
         testcommon(b"a%sb", (b'c\0d',), b'ac\0db')
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_non_ascii(self):
         testformat("\u20ac=%f", (1.0,), "\u20ac=1.000000")
 
@@ -468,8 +466,6 @@ class FormatTest(unittest.TestCase):
         self.assertIs(text % (), text)
         self.assertIs(text.format(), text)
 
-    # TODO: RustPython missing complex.__format__ implementation
-    @unittest.expectedFailure
     def test_precision(self):
         f = 1.2
         self.assertEqual(format(f, ".0f"), "1")

common/Cargo.toml

@@ -34,6 +34,7 @@ radium = { workspace = true }
 
 lock_api = "0.4"
 siphasher = "1"
+num-complex.workspace = true
 
 [target.'cfg(windows)'.dependencies]
 widestring = { workspace = true }

common/src/format.rs

@@ -1,6 +1,8 @@
 // spell-checker:ignore ddfe
 use itertools::{Itertools, PeekingNext};
+use malachite_base::num::basic::floats::PrimitiveFloat;
 use malachite_bigint::{BigInt, Sign};
+use num_complex::Complex64;
 use num_traits::FromPrimitive;
 use num_traits::{Signed, cast::ToPrimitive};
 use rustpython_literal::float;
@@ -612,6 +614,126 @@ impl FormatSpec {
         }
     }
 
+    pub fn format_complex(&self, num: &Complex64) -> Result<String, FormatSpecError> {
+        let (formatted_re, formatted_im) = self.format_complex_re_im(num)?;
+        // Enclose in parentheses if there is no format type and formatted_re is not empty
+        let magnitude_str = if self.format_type.is_none() && !formatted_re.is_empty() {
+            format!("({formatted_re}{formatted_im})")
+        } else {
+            format!("{formatted_re}{formatted_im}")
+        };
+        if let Some(FormatAlign::AfterSign) = &self.align {
+            return Err(FormatSpecError::AlignmentFlag);
+        }
+        match &self.fill.unwrap_or(' '.into()).to_char() {
+            Some('0') => Err(FormatSpecError::ZeroPadding),
+            _ => self.format_sign_and_align(&AsciiStr::new(&magnitude_str), "", FormatAlign::Right),
+        }
+    }
+
+    fn format_complex_re_im(&self, num: &Complex64) -> Result<(String, String), FormatSpecError> {
+        // Format real part
+        let mut formatted_re = String::new();
+        if num.re != 0.0 || num.re.is_negative_zero() || self.format_type.is_some() {
+            let sign_re = if num.re.is_sign_negative() && !num.is_nan() {
+                "-"
+            } else {
+                match self.sign.unwrap_or(FormatSign::Minus) {
+                    FormatSign::Plus => "+",
+                    FormatSign::Minus => "",
+                    FormatSign::MinusOrSpace => " ",
+                }
+            };
+            let re = self.format_complex_float(num.re)?;
+            formatted_re = format!("{sign_re}{re}");
+        }
+        // Format imaginary part
+        let sign_im = if num.im.is_sign_negative() && !num.im.is_nan() {
+            "-"
+        } else if formatted_re.is_empty() {
+            ""
+        } else {
+            "+"
+        };
+        let im = self.format_complex_float(num.im)?;
+        Ok((formatted_re, format!("{sign_im}{im}j")))
+    }
+
+    fn format_complex_float(&self, num: f64) -> Result<String, FormatSpecError> {
+        self.validate_format(FormatType::FixedPoint(Case::Lower))?;
+        let precision = self.precision.unwrap_or(6);
+        let magnitude = num.abs();
+        let magnitude_str = match &self.format_type {
+            Some(FormatType::Decimal)
+            | Some(FormatType::Binary)
+            | Some(FormatType::Octal)
+            | Some(FormatType::Hex(_))
+            | Some(FormatType::String)
+            | Some(FormatType::Character)
+            | Some(FormatType::Number(Case::Upper))
+            | Some(FormatType::Percentage) => {
+                let ch = char::from(self.format_type.as_ref().unwrap());
+                Err(FormatSpecError::UnknownFormatCode(ch, "complex"))
+            }
+            Some(FormatType::FixedPoint(case)) => Ok(float::format_fixed(
+                precision,
+                magnitude,
+                *case,
+                self.alternate_form,
+            )),
+            Some(FormatType::GeneralFormat(case)) | Some(FormatType::Number(case)) => {
+                let precision = if precision == 0 { 1 } else { precision };
+                Ok(float::format_general(
+                    precision,
+                    magnitude,
+                    *case,
+                    self.alternate_form,
+                    false,
+                ))
+            }
+            Some(FormatType::Exponent(case)) => Ok(float::format_exponent(
+                precision,
+                magnitude,
+                *case,
+                self.alternate_form,
+            )),
+            None => match magnitude {
+                magnitude if magnitude.is_nan() => Ok("nan".to_owned()),
+                magnitude if magnitude.is_infinite() => Ok("inf".to_owned()),
+                _ => match self.precision {
+                    Some(precision) => Ok(float::format_general(
+                        precision,
+                        magnitude,
+                        Case::Lower,
+                        self.alternate_form,
+                        true,
+                    )),
+                    None => {
+                        if magnitude.fract() == 0.0 {
+                            Ok(magnitude.trunc().to_string())
+                        } else {
+                            Ok(magnitude.to_string())
+                        }
+                    }
+                },
+            },
+        }?;
+        match &self.grouping_option {
+            Some(fg) => {
+                let sep = match fg {
+                    FormatGrouping::Comma => ',',
+                    FormatGrouping::Underscore => '_',
+                };
+                let inter = self.get_separator_interval().try_into().unwrap();
+                let len = magnitude_str.len() as i32;
+                let separated_magnitude =
+                    FormatSpec::add_magnitude_separators_for_char(magnitude_str, inter, sep, len);
+                Ok(separated_magnitude)
+            }
+            None => Ok(magnitude_str),
+        }
+    }
+
     fn format_sign_and_align<T>(
         &self,
         magnitude_str: &T,
@@ -701,6 +823,8 @@ pub enum FormatSpecError {
     NotAllowed(&'static str),
     UnableToConvert,
     CodeNotInRange,
+    ZeroPadding,
+    AlignmentFlag,
     NotImplemented(char, &'static str),
 }
 

extra_tests/snippets/builtin_format.py

@@ -165,6 +165,23 @@ assert f"{3.1415:#.1}" == "3.e+00"
 assert f"{3.1415:#.2}" == "3.1"
 assert f"{3.1415:#.3}" == "3.14"
 assert f"{3.1415:#.4}" == "3.142"
+assert f"{12.34 + 5.6j}" == "(12.34+5.6j)"
+assert f"{12.34 - 5.6j: }" == "( 12.34-5.6j)"
+assert f"{12.34 + 5.6j:20}" == "        (12.34+5.6j)"
+assert f"{12.34 + 5.6j:<20}" == "(12.34+5.6j)        "
+assert f"{-12.34 + 5.6j:^20}" == "   (-12.34+5.6j)    "
+assert f"{12.34 + 5.6j:^+20}" == "   (+12.34+5.6j)    "
+assert f"{12.34 + 5.6j:_^+20}" == "___(+12.34+5.6j)____"
+assert f"{-12.34 + 5.6j:f}" == "-12.340000+5.600000j"
+assert f"{12.34 + 5.6j:.3f}" == "12.340+5.600j"
+assert f"{12.34 + 5.6j:<30.8f}" == "12.34000000+5.60000000j       "
+assert f"{12.34 + 5.6j:g}" == "12.34+5.6j"
+assert f"{12.34 + 5.6j:e}" == "1.234000e+01+5.600000e+00j"
+assert f"{12.34 + 5.6j:E}" == "1.234000E+01+5.600000E+00j"
+assert f"{12.34 + 5.6j:^30E}" == "  1.234000E+01+5.600000E+00j  "
+assert f"{12345.6 + 7890.1j:,}" == "(12,345.6+7,890.1j)"
+assert f"{12345.6 + 7890.1j:_.3f}" == "12_345.600+7_890.100j"
+assert f"{12345.6 + 7890.1j:>+30,f}" == "  +12,345.600000+7,890.100000j"
 
 # test issue 4558
 x = 123456789012345678901234567890

vm/src/builtins/complex.rs

@@ -1,8 +1,10 @@
 use super::{PyStr, PyType, PyTypeRef, float};
 use crate::{
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
+    builtins::PyStrRef,
     class::PyClassImpl,
-    convert::{ToPyObject, ToPyResult},
+    common::format::FormatSpec,
+    convert::{IntoPyException, ToPyObject, ToPyResult},
     function::{
         OptionalArg, OptionalOption,
         PyArithmeticValue::{self, *},
@@ -388,6 +390,13 @@ impl PyComplex {
         let Complex64 { re, im } = self.value;
         (re, im)
     }
+
+    #[pymethod]
+    fn __format__(&self, spec: PyStrRef, vm: &VirtualMachine) -> PyResult<String> {
+        FormatSpec::parse(spec.as_str())
+            .and_then(|format_spec| format_spec.format_complex(&self.value))
+            .map_err(|err| err.into_pyexception(vm))
+    }
 }
 
 #[pyclass]

vm/src/format.rs

@@ -34,6 +34,12 @@ impl IntoPyException for FormatSpecError {
             }
             Self::UnableToConvert => vm.new_value_error("Unable to convert int to float"),
             Self::CodeNotInRange => vm.new_overflow_error("%c arg not in range(0x110000)"),
+            Self::ZeroPadding => {
+                vm.new_value_error("Zero padding is not allowed in complex format specifier")
+            }
+            Self::AlignmentFlag => {
+                vm.new_value_error("'=' alignment flag is not allowed in complex format specifier")
+            }
             Self::NotImplemented(c, s) => {
                 let msg = format!("Format code '{c}' for object of type '{s}' not implemented yet");
                 vm.new_value_error(msg)