Skip to main content

rustpython_codegen/
lib.rs

1//! Compile a Python AST or source code into bytecode consumable by RustPython.
2#![cfg_attr(not(feature = "std"), no_std)]
3#![doc(html_logo_url = "https://raw.githubusercontent.com/RustPython/RustPython/main/logo.png")]
4#![doc(html_root_url = "https://docs.rs/rustpython-compiler/")]
5
6#[macro_use]
7extern crate log;
8
9extern crate alloc;
10
11use alloc::{string::String, vec::Vec};
12use rustpython_compiler_core::bytecode::ConstantData;
13
14type IndexMap<K, V> = indexmap::IndexMap<K, V, rapidhash::quality::RandomState>;
15type IndexSet<T> = indexmap::IndexSet<T, rapidhash::quality::RandomState>;
16
17pub mod compile;
18pub mod error;
19pub mod ir;
20pub mod preprocess;
21mod string_parser;
22pub mod symboltable;
23mod unparse;
24
25pub use compile::CompileOpts;
26use ruff_python_ast as ast;
27use ruff_text_size::{Ranged, TextRange, TextSize, TextSlice};
28use rustpython_compiler_core::SourceFile;
29use rustpython_wtf8::Wtf8Buf;
30
31pub(crate) use compile::InternalResult;
32
33#[cfg(test)]
34pub(crate) fn constant_data_to_ast_constant_value(value: ConstantData) -> ast::ConstantValue {
35    match value {
36        ConstantData::None => ast::ConstantValue::None,
37        ConstantData::Boolean { value } => ast::ConstantValue::Boolean(value),
38        ConstantData::Str { value } => ast::ConstantValue::Str(value.to_string().into_boxed_str()),
39        ConstantData::Bytes { value } => ast::ConstantValue::Bytes(value.into_boxed_slice()),
40        ConstantData::Integer { value } => ast::ConstantValue::Integer(value.to_string().into()),
41        ConstantData::Tuple { elements } => ast::ConstantValue::Tuple(
42            elements
43                .into_iter()
44                .map(constant_data_to_ast_constant_value)
45                .collect(),
46        ),
47        ConstantData::Frozenset { elements } => ast::ConstantValue::Frozenset(
48            elements
49                .into_iter()
50                .map(constant_data_to_ast_constant_value)
51                .collect(),
52        ),
53        ConstantData::Float { value } => ast::ConstantValue::Float(value),
54        ConstantData::Complex { value } => ast::ConstantValue::Complex {
55            real: value.re,
56            imag: value.im,
57        },
58        ConstantData::Ellipsis => ast::ConstantValue::Ellipsis,
59        ConstantData::Code { .. } | ConstantData::Slice { .. } => {
60            unreachable!("ast.Constant values cannot contain code objects or slices")
61        }
62    }
63}
64
65pub(crate) fn ast_constant_value_to_constant_data(value: ast::ConstantValue) -> ConstantData {
66    match value {
67        ast::ConstantValue::None => ConstantData::None,
68        ast::ConstantValue::Boolean(value) => ConstantData::Boolean { value },
69        ast::ConstantValue::Str(value) => ConstantData::Str {
70            value: value.to_string().into(),
71        },
72        ast::ConstantValue::Bytes(value) => ConstantData::Bytes {
73            value: value.into_vec(),
74        },
75        ast::ConstantValue::Integer(value) => ConstantData::Integer {
76            value: value
77                .parse()
78                .expect("RustPython ast.Constant integer values are decimal integers"),
79        },
80        ast::ConstantValue::Tuple(elements) => ConstantData::Tuple {
81            elements: elements
82                .into_iter()
83                .map(ast_constant_value_to_constant_data)
84                .collect(),
85        },
86        ast::ConstantValue::Frozenset(elements) => ConstantData::Frozenset {
87            elements: elements
88                .into_iter()
89                .map(ast_constant_value_to_constant_data)
90                .collect(),
91        },
92        ast::ConstantValue::Float(value) => ConstantData::Float { value },
93        ast::ConstantValue::Complex { real, imag } => ConstantData::Complex {
94            value: num_complex::Complex::new(real, imag),
95        },
96        ast::ConstantValue::Ellipsis => ConstantData::Ellipsis,
97    }
98}
99
100/// The range a decorated definition covers from its `def`/`class` keyword on.
101///
102/// The parser hands over a statement that starts at the first decorator, while a
103/// definition is located at the keyword that introduces it.
104#[must_use]
105pub fn decorated_definition_range(
106    source_file: &SourceFile,
107    statement_range: TextRange,
108    decorator_list: &[ast::Decorator],
109    keyword: &str,
110) -> TextRange {
111    let Some(last_decorator) = decorator_list.last() else {
112        return statement_range;
113    };
114    let search_start = last_decorator.expression.range().end();
115    if search_start >= statement_range.end() {
116        return statement_range;
117    }
118    let search_range = TextRange::new(search_start, statement_range.end());
119    let source = source_file.source_text().slice(search_range);
120    let Some(keyword_offset) = source.find(keyword) else {
121        return statement_range;
122    };
123    let Ok(keyword_offset) = u32::try_from(keyword_offset) else {
124        return statement_range;
125    };
126    TextRange::new(
127        search_start + TextSize::new(keyword_offset),
128        statement_range.end(),
129    )
130}
131
132/// The value of a string literal.
133///
134/// The parser has nowhere to keep a lone surrogate in the value it hands over,
135/// and leaves a replacement character standing in for one, so a literal that
136/// carries any is read again from the source it was written in.
137#[must_use]
138pub fn string_literal_value(source_file: &SourceFile, string: &ast::StringLiteralValue) -> Wtf8Buf {
139    let value = string.to_str();
140    if value.contains(char::REPLACEMENT_CHARACTER) {
141        string
142            .iter()
143            .map(|part| string_literal_part_value(source_file, part))
144            .collect()
145    } else {
146        value.into()
147    }
148}
149
150/// The value of one of the literals a string literal is written as.
151#[must_use]
152pub fn string_literal_part_value(source_file: &SourceFile, string: &ast::StringLiteral) -> Wtf8Buf {
153    if string.value.contains(char::REPLACEMENT_CHARACTER) {
154        let source = source_file.source_text().slice(string.range);
155        string_parser::parse_string_literal(source, string.flags.into()).into()
156    } else {
157        string.value.to_string().into()
158    }
159}
160
161/// The value of a literal written between the interpolations of an f-string or
162/// a t-string.
163#[must_use]
164pub fn interpolated_string_literal_value(
165    source_file: &SourceFile,
166    element: &ast::InterpolatedStringLiteralElement,
167    flags: ast::AnyStringFlags,
168) -> Wtf8Buf {
169    if element.value.contains(char::REPLACEMENT_CHARACTER) {
170        let source = source_file.source_text().slice(element.range);
171        string_parser::parse_fstring_literal_element(source.into(), flags).into()
172    } else {
173        element.value.to_string().into()
174    }
175}
176
177/// The text a `{expr=}` interpolation puts in front of its value, and the range
178/// that text covers.
179///
180/// The text is the expression as it is written plus whatever the parser found
181/// around it inside the braces, with comments taken out. The range covers that
182/// surrounding text as it stands, comments included.
183#[must_use]
184pub fn interpolation_debug_text(
185    source_file: &SourceFile,
186    debug_text: &ast::DebugText,
187    expression_range: TextRange,
188) -> (String, TextRange) {
189    let leading = debug_text.leading();
190    let trailing = debug_text.trailing();
191    let text = [
192        strip_python_comments(leading).as_str(),
193        source_file.source_text().slice(expression_range),
194        strip_python_comments(trailing).as_str(),
195    ]
196    .concat();
197    let width =
198        |len: usize| TextSize::new(u32::try_from(len).expect("debug interpolation text too long"));
199    let range = TextRange::new(
200        expression_range.start() - width(leading.len()),
201        expression_range.end() + width(trailing.len()),
202    );
203    (text, range)
204}
205
206fn strip_python_comments(text: &str) -> String {
207    let chars = text.chars().collect::<Vec<_>>();
208    let mut result = String::with_capacity(text.len());
209    let mut quote = None;
210    let mut triple_quoted = false;
211    let mut escaped = false;
212    let mut in_comment = false;
213    let mut index = 0;
214
215    while index < chars.len() {
216        let ch = chars[index];
217        if in_comment {
218            if matches!(ch, '\n' | '\r') {
219                in_comment = false;
220                result.push(ch);
221            }
222            index += 1;
223            continue;
224        }
225
226        if let Some(delimiter) = quote {
227            result.push(ch);
228            if escaped {
229                escaped = false;
230            } else if ch == '\\' {
231                escaped = true;
232            } else if triple_quoted
233                && ch == delimiter
234                && chars.get(index + 1) == Some(&delimiter)
235                && chars.get(index + 2) == Some(&delimiter)
236            {
237                result.push(delimiter);
238                result.push(delimiter);
239                quote = None;
240                index += 2;
241            } else if !triple_quoted && ch == delimiter {
242                quote = None;
243            }
244            index += 1;
245            continue;
246        }
247
248        match ch {
249            '#' => in_comment = true,
250            '\'' | '"' => {
251                quote = Some(ch);
252                triple_quoted =
253                    chars.get(index + 1) == Some(&ch) && chars.get(index + 2) == Some(&ch);
254                result.push(ch);
255                if triple_quoted {
256                    result.push(ch);
257                    result.push(ch);
258                    index += 2;
259                }
260            }
261            _ => result.push(ch),
262        }
263        index += 1;
264    }
265    result
266}
267
268pub trait ToPythonName {
269    /// Returns a short name for the node suitable for use in error messages.
270    fn python_name(&self) -> &'static str;
271}
272
273impl ToPythonName for ast::Expr {
274    fn python_name(&self) -> &'static str {
275        match self {
276            Self::BoolOp { .. } | Self::BinOp { .. } | Self::UnaryOp { .. } => "operator",
277            Self::Subscript { .. } => "subscript",
278            Self::Await { .. } => "await expression",
279            Self::Yield { .. } | Self::YieldFrom { .. } => "yield expression",
280            Self::Compare { .. } => "comparison",
281            Self::Attribute { .. } => "attribute",
282            Self::Call { .. } => "function call",
283            Self::BooleanLiteral(b) => {
284                if b.value {
285                    "True"
286                } else {
287                    "False"
288                }
289            }
290            Self::EllipsisLiteral(_) => "ellipsis",
291            Self::NoneLiteral(_) => "None",
292            Self::Constant(expr) => match &expr.value {
293                ast::ConstantValue::None => "None",
294                ast::ConstantValue::Boolean(true) => "True",
295                ast::ConstantValue::Boolean(false) => "False",
296                ast::ConstantValue::Ellipsis => "ellipsis",
297                ast::ConstantValue::Tuple(_) => "tuple",
298                ast::ConstantValue::Frozenset(_) => "literal",
299                ast::ConstantValue::Str(_)
300                | ast::ConstantValue::Bytes(_)
301                | ast::ConstantValue::Integer(_)
302                | ast::ConstantValue::Float(_)
303                | ast::ConstantValue::Complex { .. } => "literal",
304            },
305            Self::NumberLiteral(_) | Self::BytesLiteral(_) | Self::StringLiteral(_) => "literal",
306            Self::Tuple(_) => "tuple",
307            Self::List { .. } => "list",
308            Self::Dict { .. } => "dict display",
309            Self::Set { .. } => "set display",
310            Self::ListComp { .. } => "list comprehension",
311            Self::DictComp { .. } => "dict comprehension",
312            Self::SetComp { .. } => "set comprehension",
313            Self::Generator { .. } => "generator expression",
314            Self::Starred { .. } => "starred",
315            Self::Slice { .. } => "slice",
316            Self::FString { .. } => "f-string expression",
317            Self::TString { .. } => "t-string expression",
318            Self::Name { .. } => "name",
319            Self::Lambda { .. } => "lambda",
320            Self::If { .. } => "conditional expression",
321            Self::Named { .. } => "named expression",
322            Self::IpyEscapeCommand(_) => "expression",
323        }
324    }
325}