Implement PyComplex's __format__ function

youknowone · Jul 4, 2025 · Jul 3, 2025 · Jul 3, 2025 · Jul 3, 2025 · Jul 4, 2025
commit 9843aecfc8e325c5e026646856f251f1617626cf
diff --git a/common/src/format.rs b/common/src/format.rs
@@ -1,6 +1,7 @@
 // spell-checker:ignore ddfe
 use itertools::{Itertools, PeekingNext};
 use malachite_bigint::{BigInt, Sign};
+use num_complex::Complex64;
 use num_traits::FromPrimitive;
 use num_traits::{Signed, cast::ToPrimitive};
 use rustpython_literal::float;
@@ -617,6 +618,105 @@ impl FormatSpec {
         }
     }
 
+    pub fn format_complex(&self, num: &Complex64) -> Result<String, FormatSpecError> {
+        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 sign_im = if num.im.is_sign_negative() && !num.im.is_nan() {
+            "-"
+        } else {
+            "+"
+        };
+        let re = self.format_complex_float(num.re)?;
+        let im = self.format_complex_float(num.im)?;
+        // Enclose in parentheses if there is no format type
+        let magnitude_str = if self.format_type.is_none() {
+            format!("({sign_re}{re}{sign_im}{im}j)")
+        } else {
+            format!("{sign_re}{re}{sign_im}{im}j")
+        };
+        if let Some('0') = &self.fill.unwrap_or(' '.into()).to_char() {
+            Err(FormatSpecError::ZeroPaddingNotAllowed)
+        } else {
+            self.format_sign_and_align(&AsciiStr::new(&magnitude_str), "", FormatAlign::Right)
+        }
+    }
+
+    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 => Ok(float::to_string(magnitude)),
+                },
+            },
+        }?;
+        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,
@@ -707,6 +807,7 @@ pub enum FormatSpecError {
     NotAllowed(&'static str),
     UnableToConvert,
     CodeNotInRange,
+    ZeroPaddingNotAllowed,
     NotImplemented(char, &'static str),
 }
 

diff --git 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<String> {
+        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
@@ -40,6 +40,9 @@ impl IntoPyException for FormatSpecError {
             FormatSpecError::CodeNotInRange => {
                 vm.new_overflow_error("%c arg not in range(0x110000)")
             }
+            FormatSpecError::ZeroPaddingNotAllowed => {
+                vm.new_value_error("Zero padding is not allowed in complex format specifier")
+            }
             FormatSpecError::NotImplemented(c, s) => {
                 let msg = format!("Format code '{c}' for object of type '{s}' not implemented yet");
                 vm.new_value_error(msg)