Skip to content

Commit d2a2120

Browse files
coolreader18coolreader18
committed
Add rustpython_literal::complex
1 parent 3ff40d2 commit d2a2120

File tree

5 files changed

+79
-75
lines changed

5 files changed

+79
-75
lines changed

‎compiler/codegen/src/unparse.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -366,7 +366,7 @@ impl<'a, 'b, 'c> Unparser<'a, 'b, 'c> {
366366
}
367367
}
368368
&ruff::Number::Complex { real, imag } => self
369-
.p(&rustpython_literal::float::complex_to_string(real, imag)
369+
.p(&rustpython_literal::complex::to_string(real, imag)
370370
.replace("inf", inf_str))?,
371371
}
372372
}

‎compiler/literal/src/complex.rs

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
use crate::float;
2+
3+
/// Convert a complex number to a string.
4+
pub fn to_string(re: f64, im: f64) -> String {
5+
// integer => drop ., fractional => float_ops
6+
let mut im_part = if im.fract() == 0.0 {
7+
im.to_string()
8+
} else {
9+
float::to_string(im)
10+
};
11+
im_part.push('j');
12+
13+
// positive empty => return im_part, integer => drop ., fractional => float_ops
14+
let re_part = if re == 0.0 {
15+
if re.is_sign_positive() {
16+
return im_part;
17+
} else {
18+
re.to_string()
19+
}
20+
} else if re.fract() == 0.0 {
21+
re.to_string()
22+
} else {
23+
float::to_string(re)
24+
};
25+
let mut result =
26+
String::with_capacity(re_part.len() + im_part.len() + 2 + im.is_sign_positive() as usize);
27+
result.push('(');
28+
result.push_str(&re_part);
29+
if im.is_sign_positive() || im.is_nan() {
30+
result.push('+');
31+
}
32+
result.push_str(&im_part);
33+
result.push(')');
34+
result
35+
}
36+
37+
/// Parse a complex number from a string.
38+
///
39+
/// Returns `Some((re, im))` on success.
40+
pub fn parse_str(s: &str) -> Option<(f64, f64)> {
41+
let s = s.trim();
42+
// Handle parentheses
43+
let s = match s.strip_prefix('(') {
44+
None => s,
45+
Some(s) => s.strip_suffix(')')?.trim(),
46+
};
47+
48+
let value = match s.strip_suffix(|c| c == 'j' || c == 'J') {
49+
None => (float::parse_str(s)?, 0.0),
50+
Some(mut s) => {
51+
let mut real = 0.0;
52+
// Find the central +/- operator. If it exists, parse the real part.
53+
for (i, w) in s.as_bytes().windows(2).enumerate() {
54+
if (w[1] == b'+' || w[1] == b'-') && !(w[0] == b'e' || w[0] == b'E') {
55+
real = float::parse_str(&s[..=i])?;
56+
s = &s[i + 1..];
57+
break;
58+
}
59+
}
60+
61+
let imag = match s {
62+
// "j", "+j"
63+
"" | "+" => 1.0,
64+
// "-j"
65+
"-" => -1.0,
66+
s => float::parse_str(s)?,
67+
};
68+
69+
(real, imag)
70+
}
71+
};
72+
Some(value)
73+
}

‎compiler/literal/src/float.rs

Lines changed: 0 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -178,39 +178,6 @@ pub fn to_string(value: f64) -> String {
178178
}
179179
}
180180

181-
pub fn complex_to_string(re: f64, im: f64) -> String {
182-
// integer => drop ., fractional => float_ops
183-
let mut im_part = if im.fract() == 0.0 {
184-
im.to_string()
185-
} else {
186-
to_string(im)
187-
};
188-
im_part.push('j');
189-
190-
// positive empty => return im_part, integer => drop ., fractional => float_ops
191-
let re_part = if re == 0.0 {
192-
if re.is_sign_positive() {
193-
return im_part;
194-
} else {
195-
re.to_string()
196-
}
197-
} else if re.fract() == 0.0 {
198-
re.to_string()
199-
} else {
200-
to_string(re)
201-
};
202-
let mut result =
203-
String::with_capacity(re_part.len() + im_part.len() + 2 + im.is_sign_positive() as usize);
204-
result.push('(');
205-
result.push_str(&re_part);
206-
if im.is_sign_positive() || im.is_nan() {
207-
result.push('+');
208-
}
209-
result.push_str(&im_part);
210-
result.push(')');
211-
result
212-
}
213-
214181
pub fn from_hex(s: &str) -> Option<f64> {
215182
if let Ok(f) = hexf_parse::parse_hexf64(s, false) {
216183
return Some(f);

‎compiler/literal/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
pub mod char;
2+
pub mod complex;
23
pub mod escape;
34
pub mod float;
45
pub mod format;

‎vm/src/builtins/complex.rs

Lines changed: 4 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -179,13 +179,13 @@ impl Constructor for PyComplex {
179179
"complex() can't take second arg if first is a string".to_owned(),
180180
));
181181
}
182-
let value = s
182+
let (re, im) = s
183183
.to_str()
184-
.and_then(|s| parse_str(s.trim()))
184+
.and_then(rustpython_literal::complex::parse_str)
185185
.ok_or_else(|| {
186186
vm.new_value_error("complex() arg is a malformed string".to_owned())
187187
})?;
188-
return Self::from(value)
188+
return Self::from(Complex64 { re, im })
189189
.into_ref_with_type(vm, cls)
190190
.map(Into::into);
191191
} else {
@@ -494,7 +494,7 @@ impl Representable for PyComplex {
494494
// TODO: when you fix this, move it to rustpython_common::complex::repr and update
495495
// ast/src/unparse.rs + impl Display for Constant in ast/src/constant.rs
496496
let Complex64 { re, im } = zelf.value;
497-
Ok(rustpython_literal::float::complex_to_string(re, im))
497+
Ok(rustpython_literal::complex::to_string(re, im))
498498
}
499499
}
500500

@@ -519,40 +519,3 @@ pub struct ComplexArgs {
519519
#[pyarg(any, optional)]
520520
imag: OptionalArg<PyObjectRef>,
521521
}
522-
523-
fn parse_str(s: &str) -> Option<Complex64> {
524-
// Handle parentheses
525-
let s = match s.strip_prefix('(') {
526-
None => s,
527-
Some(s) => match s.strip_suffix(')') {
528-
None => return None,
529-
Some(s) => s.trim(),
530-
},
531-
};
532-
533-
let value = match s.strip_suffix(|c| c == 'j' || c == 'J') {
534-
None => Complex64::new(crate::literal::float::parse_str(s)?, 0.0),
535-
Some(mut s) => {
536-
let mut real = 0.0;
537-
// Find the central +/- operator. If it exists, parse the real part.
538-
for (i, w) in s.as_bytes().windows(2).enumerate() {
539-
if (w[1] == b'+' || w[1] == b'-') && !(w[0] == b'e' || w[0] == b'E') {
540-
real = crate::literal::float::parse_str(&s[..=i])?;
541-
s = &s[i + 1..];
542-
break;
543-
}
544-
}
545-
546-
let imag = match s {
547-
// "j", "+j"
548-
"" | "+" => 1.0,
549-
// "-j"
550-
"-" => -1.0,
551-
s => crate::literal::float::parse_str(s)?,
552-
};
553-
554-
Complex64::new(real, imag)
555-
}
556-
};
557-
Some(value)
558-
}

0 commit comments

Comments
 (0)