rustpython_literal/
complex.rs

1use crate::float;
2use alloc::borrow::ToOwned;
3use alloc::string::{String, ToString};
4
5/// Format a single complex component (real or imag) for `repr`.
6/// Uses scientific notation when `|value| < 1e-4` or `|value| >= 1e16`
7/// (matching CPython's `PyOS_double_to_string(format='r')`), otherwise
8/// Rust's default `Display`, which drops the trailing `.0` for
9/// integer-valued floats.
10///
11/// This differs from `float::to_string` only in that integer values in
12/// the normal range render as `"1"` rather than `"1.0"` — complex repr
13/// formats `1+2j` as `"(1+2j)"`, not `"(1.0+2.0j)"`.
14fn component_to_string(value: f64) -> String {
15    let lit = alloc::format!("{value:e}");
16    if let Some(position) = lit.find('e') {
17        let significand = &lit[..position];
18        let exponent = lit[position + 1..].parse::<i32>().unwrap();
19        if exponent < 16 && exponent > -5 {
20            // Normal magnitude — Rust's default Display emits "1" for 1.0,
21            // "1.5" for 1.5, "1000000000000000" for 1e15, etc.
22            float::prefer_cpython_tie_repr(value.to_string(), value)
23        } else {
24            float::prefer_cpython_tie_repr(alloc::format!("{significand}e{exponent:+#03}"), value)
25        }
26    } else {
27        // nan / inf / -inf — `format!("{x:e}")` produces e.g. "NaN" with no
28        // exponent marker; lowercase to match Python.
29        let mut s = value.to_string();
30        s.make_ascii_lowercase();
31        s
32    }
33}
34
35/// Convert a complex number to a string.
36#[must_use]
37pub fn to_string(re: f64, im: f64) -> String {
38    let mut im_part = component_to_string(im);
39    im_part.push('j');
40
41    // positive empty => return im_part, integer => drop ., fractional => float_ops
42    let re_part = if re == 0.0 {
43        if re.is_sign_positive() {
44            return im_part;
45        }
46        "-0".to_owned()
47    } else {
48        component_to_string(re)
49    };
50    let mut result =
51        String::with_capacity(re_part.len() + im_part.len() + 2 + im.is_sign_positive() as usize);
52    result.push('(');
53    result.push_str(&re_part);
54    if im.is_sign_positive() || im.is_nan() {
55        result.push('+');
56    }
57    result.push_str(&im_part);
58    result.push(')');
59    result
60}
61
62/// Parse a complex number from a string.
63///
64/// Returns `Some((re, im))` on success.
65pub fn parse_str(s: &str) -> Option<(f64, f64)> {
66    let s = s.trim();
67    // Handle parentheses
68    let s = match s.strip_prefix('(') {
69        None => s,
70        Some(s) => s.strip_suffix(')')?.trim(),
71    };
72
73    // Whitespace is only allowed around the whole string and the optional
74    // parentheses, never inside the numeric token. Reject it here so that
75    // `float::parse_str` (which tolerates surrounding whitespace on a part)
76    // does not let e.g. "1 +2j" through.
77    if s.contains(char::is_whitespace) {
78        return None;
79    }
80
81    let value = match s.strip_suffix(|c| c == 'j' || c == 'J') {
82        None => (float::parse_str(s)?, 0.0),
83        Some(mut s) => {
84            let mut real = 0.0;
85            // Find the central +/- operator. If it exists, parse the real part.
86            for (i, w) in s.as_bytes().windows(2).enumerate() {
87                if (w[1] == b'+' || w[1] == b'-') && !(w[0] == b'e' || w[0] == b'E') {
88                    real = float::parse_str(&s[..=i])?;
89                    s = &s[i + 1..];
90                    break;
91                }
92            }
93
94            let imag = match s {
95                // "j", "+j"
96                "" | "+" => 1.0,
97                // "-j"
98                "-" => -1.0,
99                s => float::parse_str(s)?,
100            };
101
102            (real, imag)
103        }
104    };
105    Some(value)
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn parse_rejects_internal_whitespace() {
114        // Whitespace inside the numeric token is invalid, even where a bare
115        // `float::parse_str` on a fragment would tolerate it.
116        for s in [
117            "1 +2j", "1 2j", "1 +2 j", "+ 1j", "1.5 j", "(1 +2j)", "2 -3j",
118        ] {
119            assert_eq!(parse_str(s), None, "{s:?} must not parse");
120        }
121    }
122
123    #[test]
124    fn parse_allows_surrounding_and_paren_whitespace() {
125        for s in ["  1+2j  ", " (1+2j) ", "( 1+2j )"] {
126            assert_eq!(parse_str(s), Some((1.0, 2.0)), "{s:?}");
127        }
128    }
129
130    #[test]
131    fn parse_basic() {
132        assert_eq!(parse_str("1"), Some((1.0, 0.0)));
133        assert_eq!(parse_str("1j"), Some((0.0, 1.0)));
134        assert_eq!(parse_str("j"), Some((0.0, 1.0)));
135        assert_eq!(parse_str("-j"), Some((0.0, -1.0)));
136        assert_eq!(parse_str("1+2j"), Some((1.0, 2.0)));
137        assert_eq!(parse_str("1e5j"), Some((0.0, 1e5)));
138        assert_eq!(parse_str("1_000"), Some((1000.0, 0.0)));
139        assert_eq!(parse_str(""), None);
140        assert_eq!(parse_str("abc"), None);
141    }
142}
rustpython_literal/complex.rs

rustpython_literal/
complex.rs
