Skip to main content

rustpython_common/encodings/
escape.rs

1//! Bytes-to-bytes Python string-literal escape transform (`escape_encode` /
2//! `escape_decode`).
3
4use alloc::string::String;
5use alloc::vec::Vec;
6
7/// Encode `data` the way `escape_encode` / `string_escape_encode` does.
8#[must_use]
9pub fn encode(data: &[u8]) -> Vec<u8> {
10    let mut out = Vec::with_capacity(data.len());
11    for &byte in data {
12        match byte {
13            b'\t' => out.extend_from_slice(b"\\t"),
14            b'\n' => out.extend_from_slice(b"\\n"),
15            b'\r' => out.extend_from_slice(b"\\r"),
16            b'\\' => out.extend_from_slice(b"\\\\"),
17            b'\'' => out.extend_from_slice(b"\\'"),
18            0x20..=0x7e => out.push(byte),
19            value => out.extend_from_slice(format!("\\x{value:02x}").as_bytes()),
20        }
21    }
22    out
23}
24
25/// How a malformed `\x` escape is reported.
26#[derive(Clone, Copy, Debug, PartialEq, Eq)]
27pub enum EscapeErrorMode {
28    Strict,
29    Replace,
30    Ignore,
31}
32
33impl EscapeErrorMode {
34    #[must_use]
35    pub fn from_name(name: &str) -> Option<Self> {
36        match name {
37            "strict" => Some(Self::Strict),
38            "replace" => Some(Self::Replace),
39            "ignore" => Some(Self::Ignore),
40            _ => None,
41        }
42    }
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub enum EscapeDecodeError {
47    TrailingBackslash,
48    InvalidHex { position: usize },
49    UnknownHandler { name: String },
50}
51
52/// Decode a Python bytes-literal escape sequence.
53///
54/// `warning` is the first invalid-escape deprecation text, if any.
55pub fn decode(
56    data: &[u8],
57    errors: EscapeErrorMode,
58) -> Result<(Vec<u8>, Option<String>), EscapeDecodeError> {
59    let mut out = Vec::with_capacity(data.len());
60    let mut pos = 0usize;
61    let mut warning = None;
62    while pos < data.len() {
63        if data[pos] != b'\\' {
64            out.push(data[pos]);
65            pos += 1;
66            continue;
67        }
68        let escape_start = pos;
69        pos += 1;
70        if pos == data.len() {
71            return Err(EscapeDecodeError::TrailingBackslash);
72        }
73        let ch = data[pos];
74        pos += 1;
75        match ch {
76            b'\n' => {}
77            b'\\' => out.push(b'\\'),
78            b'\'' => out.push(b'\''),
79            b'"' => out.push(b'"'),
80            b'b' => out.push(0x08),
81            b'f' => out.push(0x0c),
82            b't' => out.push(b'\t'),
83            b'n' => out.push(b'\n'),
84            b'r' => out.push(b'\r'),
85            b'v' => out.push(0x0b),
86            b'a' => out.push(0x07),
87            b'0'..=b'7' => {
88                let octal_start = pos - 1;
89                while pos < data.len()
90                    && pos < octal_start + 3
91                    && (b'0'..=b'7').contains(&data[pos])
92                {
93                    pos += 1;
94                }
95                let raw = data[octal_start..pos]
96                    .iter()
97                    .fold(0u16, |value, digit| value * 8 + u16::from(digit - b'0'));
98                if raw >= 256 && warning.is_none() {
99                    warning = Some(invalid_escape_warning(
100                        "b",
101                        &String::from_utf8_lossy(&data[octal_start..pos]),
102                        true,
103                    ));
104                }
105                out.push(raw as u8);
106            }
107            b'x' => {
108                let hi = data.get(pos).and_then(|byte| (*byte as char).to_digit(16));
109                let lo = data
110                    .get(pos + 1)
111                    .and_then(|byte| (*byte as char).to_digit(16));
112                if let (Some(hi), Some(lo)) = (hi, lo) {
113                    out.push((hi * 16 + lo) as u8);
114                    pos += 2;
115                } else {
116                    match errors {
117                        EscapeErrorMode::Strict => {
118                            return Err(EscapeDecodeError::InvalidHex {
119                                position: escape_start,
120                            });
121                        }
122                        EscapeErrorMode::Replace => out.push(b'?'),
123                        EscapeErrorMode::Ignore => {}
124                    }
125                    if data.get(pos).is_some_and(u8::is_ascii_hexdigit) {
126                        pos += 1;
127                    }
128                }
129            }
130            other => {
131                out.push(b'\\');
132                pos -= 1;
133                if warning.is_none() {
134                    warning = Some(invalid_escape_warning(
135                        "b",
136                        &char::from(other).to_string(),
137                        false,
138                    ));
139                }
140            }
141        }
142    }
143    Ok((out, warning))
144}
145
146fn invalid_escape_warning(prefix: &str, sequence: &str, octal: bool) -> String {
147    let kind = if octal {
148        "an invalid octal escape sequence"
149    } else {
150        "an invalid escape sequence"
151    };
152    format!("{prefix}\"\\{sequence}\" is {kind}. Such sequences will not work in the future. ")
153}