8000 Implement PyComplex's __format__ function (#5900) · RustPython/RustPython@0a59c1c · GitHub
[go: up one dir, main page]
Skip to content

Commit 0a59c1c

Browse files
yt2byt2b
authored
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
1 parent 694fe50 commit 0a59c1c

File tree

8 files changed

+159
-7
lines changed

8 files changed

+159
-7
lines changed

‎Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Lib/test/test_complex.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -719,8 +719,6 @@ def test_repr_roundtrip(self):
719719
self.assertFloatsAreIdentical(0.0 + z.imag,
720720
0.0 + roundtrip.imag)
721721

722-
# TODO: RUSTPYTHON
723-
@unittest.expectedFailure
724722
def test_format(self):
725723
# empty format string is same as str()
726724
self.assertEqual(format(1+3j, ''), str(1+3j))

‎Lib/test/test_format.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -397,8 +397,6 @@ def test_nul(self):
397397
testformat("a%sb", ('c\0d',), 'ac\0db')
398398
testcommon(b"a%sb", (b'c\0d',), b'ac\0db')
399399

400-
# TODO: RUSTPYTHON
401-
@unittest.expectedFailure
402400
def test_non_ascii(self):
403401
testformat("\u20ac=%f", (1.0,), "\u20ac=1.000000")
404402

@@ -468,8 +466,6 @@ def test_optimisations(self):
468466
self.assertIs(text % (), text)
469467
self.assertIs(text.format(), text)
470468

471-
# TODO: RustPython missing complex.__format__ implementation
472-
@unittest.expectedFailure
473469
def test_precision(self):
474470
f = 1.2
475471
self.assertEqual(format(f, ".0f"), "1")

‎common/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ radium = { workspace = true }
3434

3535
lock_api = "0.4"
3636
siphasher = "1"
37+
num-complex.workspace = true
3738

3839
[target.'cfg(windows)'.dependencies]
3940
widestring = { workspace = true }

‎common/src/format.rs

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
// spell-checker:ignore ddfe
22
use itertools::{Itertools, PeekingNext};
3+
use malachite_base::num::basic::floats::PrimitiveFloat;
34
use malachite_bigint::{BigInt, Sign};
5+
use num_complex::Complex64;
46
use num_traits::FromPrimitive;
57
use num_traits::{Signed, cast::ToPrimitive};
68
use rustpython_literal::float;
@@ -612,6 +614,126 @@ impl FormatSpec {
612614
}
613615
}
614616

617+
pub fn format_complex(&self, num: &Complex64) -> Result<String, FormatSpecError> {
618+
let (formatted_re, formatted_im) = self.format_complex_re_im(num)?;
619+
// Enclose in parentheses if there is no format type and formatted_re is not empty
620+
let magnitude_str = if self.format_type.is_none() && !formatted_re.is_empty() {
621+
format!("({formatted_re}{formatted_im})")
622+
} else {
623+
format!("{formatted_re}{formatted_im}")
624+
};
625+
if let Some(FormatAlign::AfterSign) = &self.align {
626+
return Err(FormatSpecError::AlignmentFlag);
627+
}
628+
match &self.fill.unwrap_or(' '.into()).to_char() {
629+
Some('0') => Err(FormatSpecError::ZeroPadding),
630+
_ => self.format_sign_and_align(&AsciiStr::new(&magnitude_str), "", FormatAlign::Right),
631+
}
632+
}
633+
634+
fn format_complex_re_im(&self, num: &Complex64) -> Result<(String, String), FormatSpecError> {
635+
// Format real part
636+
let mut formatted_re = String::new();
637+
if num.re != 0.0 || num.re.is_negative_zero() || self.format_type.is_some() {
638+
let sign_re = if num.re.is_sign_negative() && !num.is_nan() {
639+
"-"
640+
} else {
641+
match self.sign.unwrap_or(FormatSign::Minus) {
642+
FormatSign::Plus => "+",
643+
FormatSign::Minus => "",
644+
FormatSign::MinusOrSpace => " ",
645+
}
646+
};
647+
let re = self.format_complex_float(num.re)?;
648+
formatted_re = format!("{sign_re}{re}");
649+
}
650+
// Format imaginary part
651+
let sign_im = if num.im.is_sign_negative() && !num.im.is_nan() {
652< 10000 /td>+
"-"
653+
} else if formatted_re.is_empty() {
654+
""
655+
} else {
656+
"+"
657+
};
658+
let im = self.format_complex_float(num.im)?;
659+
Ok((formatted_re, format!("{sign_im}{im}j")))
660+
}
661+
662+
fn format_complex_float(&self, num: f64) -> Result<String, FormatSpecError> {
663+
self.validate_format(FormatType::FixedPoint(Case::Lower))?;
664+
let precision = self.precision.unwrap_or(6);
665+
let magnitude = num.abs();
666+
let magnitude_str = match &self.format_type {
667+
Some(FormatType::Decimal)
668+
| Some(FormatType::Binary)
669+
| Some(FormatType::Octal)
670+
| Some(FormatType::Hex(_))
671+
| Some(FormatType::String)
672+
| Some(FormatType::Character)
673+
| Some(FormatType::Number(Case::Upper))
674+
| Some(FormatType::Percentage) => {
675+
let ch = char::from(self.format_type.as_ref().unwrap());
676+
Err(FormatSpecError::UnknownFormatCode(ch, "complex"))
677+
}
678+
Some(FormatType::FixedPoint(case)) => Ok(float::format_fixed(
679+
precision,
680+
magnitude,
681+
*case,
682+
self.alternate_form,
683+
)),
684+
Some(FormatType::GeneralFormat(case)) | Some(FormatType::Number(case)) => {
685+
let precision = if precision == 0 { 1 } else { precision };
686+
Ok(float::format_general(
687+
precision,
688+
magnitude,
689+
*case,
690+
self.alternate_form,
691+
false,
692+
))
693+
}
694+
Some(FormatType::Exponent(case)) => Ok(float::format_exponent(
695+
precision,
696+
magnitude,
697+
*case,
698+
self.alternate_form,
699+
)),
700+
None => match magnitude {
701+
magnitude if magnitude.is_nan() => Ok("nan".to_owned()),
702+
magnitude if magnitude.is_infinite() => Ok("inf".to_owned()),
703+
_ => match self.precision {
704+
Some(precision) => Ok(float::format_general(
705+
precision,
706+
magnitude,
707+
Case::Lower,
708+
self.alternate_form,
709+
true,
710+
)),
711+
None => {
712+
if magnitude.fract() == 0.0 {
713+
Ok(magnitude.trunc().to_string())
714+
} else {
715+
Ok(magnitude.to_string())
716+
}
717+
}
718+
},
719+
},
720+
}?;
721+
match &self.grouping_option {
722+
Some(fg) => {
723+
let sep = match fg {
724+
FormatGrouping::Comma => ',',
725+
FormatGrouping::Underscore => '_',
726+
};
727+
let inter = self.get_separator_interval().try_into().unwrap();
728+
let len = magnitude_str.len() as i32;
729+
let separated_magnitude =
730+
FormatSpec::add_magnitude_separators_for_char(magnitude_str, inter, sep, len);
731+
Ok(separated_magnitude)
732+
}
733+
None => Ok(magnitude_str),
734+
}
735+
}
736+
615737
fn format_sign_and_align<T>(
616738
&self,
617739
magnitude_str: &T,
@@ -701,6 +823,8 @@ pub enum FormatSpecError {
701823
NotAllowed(&'static str),
702824
UnableToConvert,
703825
CodeNotInRange,
826+
ZeroPadding,
827+
AlignmentFlag,
704828
NotImplemented(char, &'static str),
705829
}
706830

‎extra_tests/snippets/builtin_format.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,23 @@ def test_zero_padding():
165165
assert f"{3.1415:#.2}" == "3.1"
166166
assert f"{3.1415:#.3}" == "3.14"
167 93C6 167
assert f"{3.1415:#.4}" == "3.142"
168+
assert f"{12.34 + 5.6j}" == "(12.34+5.6j)"
169+
assert f"{12.34 - 5.6j: }" == "( 12.34-5.6j)"
170+
assert f"{12.34 + 5.6j:20}" == " (12.34+5.6j)"
171+
assert f"{12.34 + 5.6j:<20}" == "(12.34+5.6j) "
172+
assert f"{-12.34 + 5.6j:^20}" == " (-12.34+5.6j) "
173+
assert f"{12.34 + 5.6j:^+20}" == " (+12.34+5.6j) "
174+
assert f"{12.34 + 5.6j:_^+20}" == "___(+12.34+5.6j)____"
175+
assert f"{-12.34 + 5.6j:f}" == "-12.340000+5.600000j"
176+
assert f"{12.34 + 5.6j:.3f}" == "12.340+5.600j"
177+
assert f"{12.34 + 5.6j:<30.8f}" == "12.34000000+5.60000000j "
178+
assert f"{12.34 + 5.6j:g}" == "12.34+5.6j"
179+
assert f"{12.34 + 5.6j:e}" == "1.234000e+01+5.600000e+00j"
180+
assert f"{12.34 + 5.6j:E}" == "1.234000E+01+5.600000E+00j"
181+
assert f"{12.34 + 5.6j:^30E}" == " 1.234000E+01+5.600000E+00j "
182+
assert f"{12345.6 + 7890.1j:,}" == "(12,345.6+7,890.1j)"
183+
assert f"{12345.6 + 7890.1j:_.3f}" == "12_345.600+7_890.100j"
184+
assert f"{12345.6 + 7890.1j:>+30,f}" == " +12,345.600000+7,890.100000j"
168185

169186
# test issue 4558
170187
x = 123456789012345678901234567890

‎vm/src/builtins/complex.rs

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
use super::{PyStr, PyType, PyTypeRef, float};
22
use crate::{
33
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
4+
builtins::PyStrRef,
45
class::PyClassImpl,
5-
convert::{ToPyObject, ToPyResult},
6+
common::format::FormatSpec,
7+
convert::{IntoPyException, ToPyObject, ToPyResult},
68
function::{
79
OptionalArg, OptionalOption,
810
PyArithmeticValue::{self, *},
@@ -388,6 +390,13 @@ impl PyComplex {
388390
let Complex64 { re, im } = self.value;
389391
(re, im)
390392
}
393+
394+
#[pymethod]
395+
fn __format__(&self, spec: PyStrRef, vm: &VirtualMachine) -> PyResult<String> {
396+
FormatSpec::parse(spec.as_str())
397+
.and_then(|format_spec| format_spec.format_complex(&self.value))
398+
.map_err(|err| err.into_pyexception(vm))
399+
}
391400
}
392401

393402
#[pyclass]

‎vm/src/format.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,12 @@ impl IntoPyException for FormatSpecError {
3434
}
3535
Self::UnableToConvert => vm.new_value_error("Unable to convert int to float"),
3636
Self::CodeNotInRange => vm.new_overflow_error("%c arg not in range(0x110000)"),
37+
Self::ZeroPadding => {
38+
vm.new_value_error("Zero padding is not allowed in complex format specifier")
39+
}
40+
Self::AlignmentFlag => {
41+
vm.new_value_error("'=' alignment flag is not allowed in complex format specifier")
42+
}
3743
Self::NotImplemented(c, s) => {
3844
let msg = format!("Format code '{c}' for object of type '{s}' not implemented yet");
3945
vm.new_value_error(msg)

0 commit comments

Comments
 (0)
0