From 0a59c1cad339a0a390e303f8146df7484e861e8e Mon Sep 17 00:00:00 2001 From: yt2b <76801443+yt2b@users.noreply.github.com> Date: Fri, 4 Jul 2025 22:35:28 +0900 Subject: [PATCH] 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 | 1 + Lib/test/test_complex.py | 2 - Lib/test/test_format.py | 4 - common/Cargo.toml | 1 + common/src/format.rs | 124 +++++++++++++++++++++++++ extra_tests/snippets/builtin_format.py | 17 ++++ vm/src/builtins/complex.rs | 11 ++- vm/src/format.rs | 6 ++ 8 files changed, 159 insertions(+), 7 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f05973255..c6dee806c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2300,6 +2300,7 @@ dependencies = [ "malachite-bigint", "malachite-q", "memchr", + "num-complex", "num-traits", "once_cell", "parking_lot", diff --git a/Lib/test/test_complex.py b/Lib/test/test_complex.py index 106182cab..dd3c4f281 100644 --- a/Lib/test/test_complex.py +++ b/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)) diff --git a/Lib/test/test_format.py b/Lib/test/test_format.py index 187270d5b..1a26d95bc 100644 --- a/Lib/test/test_format.py +++ b/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") diff --git a/common/Cargo.toml b/common/Cargo.toml index 4eab8440d..94704d18c 100644 --- a/common/Cargo.toml +++ b/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 } diff --git a/common/src/format.rs b/common/src/format.rs index feebc3893..3aee85566 100644 --- a/common/src/format.rs +++ b/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 { + 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 { + 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( &self, magnitude_str: &T, @@ -701,6 +823,8 @@ pub enum FormatSpecError { NotAllowed(&'static str), UnableToConvert, CodeNotInRange, + ZeroPadding, + AlignmentFlag, NotImplemented(char, &'static str), } diff --git a/extra_tests/snippets/builtin_format.py b/extra_tests/snippets/builtin_format.py index ac96b6545..457adaa13 100644 --- a/extra_tests/snippets/builtin_format.py +++ b/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 diff --git a/vm/src/builtins/complex.rs b/vm/src/builtins/complex.rs index 28adb7297..11c782664 100644 --- a/vm/src/builtins/complex.rs +++ b/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 { + FormatSpec::parse(spec.as_str()) + .and_then(|format_spec| format_spec.format_complex(&self.value)) + .map_err(|err| err.into_pyexception(vm)) + } } #[pyclass] diff --git a/vm/src/format.rs b/vm/src/format.rs index 3eca5ce1b..1b25c9b67 100644 --- a/vm/src/format.rs +++ b/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)