Skip to content

Commit 7333fd1

Browse files
youknowoneyouknowone
committed
Lib/ast.py
1 parent 3dbb6a6 commit 7333fd1

File tree

4 files changed

+37
-22
lines changed

4 files changed

+37
-22
lines changed

‎vm/src/builtins.rs

+13-4
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ use crate::pyobject::{
3030
};
3131
use crate::scope::Scope;
3232
use crate::vm::VirtualMachine;
33-
33+
use crate::stdlib::ast;
3434
#[cfg(not(target_arch = "wasm32"))]
3535
use crate::stdlib::io::io_open;
3636

@@ -125,7 +125,7 @@ struct CompileArgs {
125125
}
126126

127127
#[cfg(feature = "rustpython-compiler")]
128-
fn builtin_compile(args: CompileArgs, vm: &VirtualMachine) -> PyResult<PyCodeRef> {
128+
fn builtin_compile(args: CompileArgs, vm: &VirtualMachine) -> PyResult {
129129
// TODO: compile::compile should probably get bytes
130130
let source = match &args.source {
131131
Either::A(string) => string.as_str(),
@@ -138,8 +138,17 @@ fn builtin_compile(args: CompileArgs, vm: &VirtualMachine) -> PyResult<PyCodeRef
138138
.parse::<compile::Mode>()
139139
.map_err(|err| vm.new_value_error(err.to_string()))?;
140140

141-
vm.compile(source, mode, args.filename.as_str().to_string())
142-
.map_err(|err| vm.new_syntax_error(&err))
141+
let flags = args
142+
.flags
143+
.map_or(Ok(0), |v| i32::try_from_object(vm, v.into_object()))?;
144+
145+
if (flags & ast::PY_COMPILE_FLAG_AST_ONLY).is_zero() {
146+
vm.compile(&source, mode, args.filename.value.to_string())
147+
.map(|o| o.into_object())
148+
.map_err(|err| vm.new_syntax_error(&err))
149+
} else {
150+
ast::source_to_ast(&vm, &source)
151+
}
143152
}
144153

145154
fn builtin_delattr(obj: PyObjectRef, attr: PyStringRef, vm: &VirtualMachine) -> PyResult<()> {

‎vm/src/function.rs

+10
Original file line numberDiff line numberDiff line change
@@ -381,6 +381,16 @@ impl<T> OptionalArg<T> {
381381
}
382382
}
383383

384+
pub fn map_or<U, F>(self, default: U, f: F) -> U
385+
where
386+
F: FnOnce(T) -> U,
387+
{
388+
match self {
389+
Present(value) => f(value),
390+
Missing => default,
391+
}
392+
}
393+
384394
pub fn map_or_else<U, D, F>(self, default: D, f: F) -> U
385395
where
386396
D: FnOnce() -> U,

‎vm/src/stdlib/ast.rs

+12-16
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ use num_complex::Complex64;
1010
use rustpython_parser::{ast, parser};
1111

1212
use crate::obj::objlist::PyListRef;
13-
use crate::obj::objstr::PyStringRef;
1413
use crate::obj::objtype::PyClassRef;
1514
use crate::pyobject::{PyObjectRef, PyRef, PyResult, PyValue};
1615
use crate::vm::VirtualMachine;
@@ -19,9 +18,12 @@ use crate::vm::VirtualMachine;
1918
struct AstNode;
2019
type AstNodeRef = PyRef<AstNode>;
2120

21+
const MODULE_NAME: &str = "_ast";
22+
pub const PY_COMPILE_FLAG_AST_ONLY: i32 = 0x0400;
23+
2224
impl PyValue for AstNode {
2325
fn class(vm: &VirtualMachine) -> PyClassRef {
24-
vm.class("ast", "AST")
26+
vm.class(MODULE_NAME, "AST")
2527
}
2628
}
2729

@@ -48,14 +50,16 @@ macro_rules! node {
4850
}
4951
}
5052

51-
fn program_to_ast(vm: &VirtualMachine, program: &ast::Program) -> PyResult<AstNodeRef> {
52-
let py_body = statements_to_ast(vm, &program.statements)?;
53-
Ok(node!(vm, Module, { body => py_body }))
53+
pub(crate) fn source_to_ast(vm: &VirtualMachine, source: &str) -> PyResult {
54+
let internal_ast =
55+
parser::parse_program(source).map_err(|err| vm.new_value_error(format!("{}", err)))?;
56+
let py_body = statements_to_ast(vm, &internal_ast.statements)?;
57+
Ok(node!(vm, Module, { body => py_body }).into_object())
5458
}
5559

5660
// Create a node class instance
5761
fn create_node(vm: &VirtualMachine, name: &str) -> PyResult<AstNodeRef> {
58-
AstNode.into_ref_with_type(vm, vm.class("ast", name))
62+
AstNode.into_ref_with_type(vm, vm.class(MODULE_NAME, name))
5963
}
6064

6165
fn statements_to_ast(vm: &VirtualMachine, statements: &[ast::Statement]) -> PyResult<PyListRef> {
@@ -630,20 +634,11 @@ fn string_to_ast(vm: &VirtualMachine, string: &ast::StringGroup) -> PyResult<Ast
630634
Ok(string)
631635
}
632636

633-
fn ast_parse(source: PyStringRef, vm: &VirtualMachine) -> PyResult<AstNodeRef> {
634-
let internal_ast = parser::parse_program(source.as_str())
635-
.map_err(|err| vm.new_value_error(format!("{}", err)))?;
636-
// source.clone();
637-
program_to_ast(&vm, &internal_ast)
638-
}
639-
640-
#[allow(clippy::cognitive_complexity)]
641637
pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
642638
let ctx = &vm.ctx;
643639

644640
let ast_base = py_class!(ctx, "AST", ctx.object(), {});
645-
py_module!(vm, "ast", {
646-
"parse" => ctx.new_rustfunc(ast_parse),
641+
py_module!(vm, MODULE_NAME, {
647642
"AST" => ast_base.clone(),
648643
// TODO: There's got to be a better way!
649644
"alias" => py_class!(ctx, "alias", ast_base.clone(), {}),
@@ -709,5 +704,6 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
709704
"withitem" => py_class!(ctx, "withitem", ast_base.clone(), {}),
710705
"Yield" => py_class!(ctx, "Yield", ast_base.clone(), {}),
711706
"YieldFrom" => py_class!(ctx, "YieldFrom", ast_base.clone(), {}),
707+
"PyCF_ONLY_AST" => ctx.new_int(PY_COMPILE_FLAG_AST_ONLY),
712708
})
713709
}

‎vm/src/stdlib/mod.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
pub mod array;
22
#[cfg(feature = "rustpython-parser")]
3-
mod ast;
3+
pub(crate) mod ast;
44
mod binascii;
55
mod codecs;
66
mod collections;
@@ -86,7 +86,7 @@ pub fn get_module_inits() -> HashMap<String, StdlibInitFunc> {
8686
#[cfg(feature = "rustpython-parser")]
8787
{
8888
modules.insert(
89-
"ast".to_string(),
89+
"_ast".to_string(),
9090
Box::new(ast::make_module) as StdlibInitFunc,
9191
);
9292
modules.insert("keyword".to_string(), Box::new(keyword::make_module));

0 commit comments

Comments
 (0)