rustpython_literal/
complex.rs

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

rustpython_literal/
complex.rs
