Skip to content

Refactor _ast module for clarity #7031

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
youknowone merged 1 commit into from
Feb 6, 2026
Merged

Refactor _ast module for clarity #7031

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 15 additions & 3 deletions crates/vm/src/stdlib/ast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,18 @@ mod string;
mod type_ignore;
mod type_parameters;

/// Return the cached singleton instance for an operator/context node type,
/// or create a new instance if none exists.
fn singleton_node_to_object(vm: &VirtualMachine, node_type: &'static Py<PyType>) -> PyObjectRef {
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
return instance;
}
NodeAst
.into_ref_with_type(vm, node_type.to_owned())
.unwrap()
.into()
}

fn get_node_field(vm: &VirtualMachine, obj: &PyObject, field: &'static str, typ: &str) -> PyResult {
vm.get_attribute_opt(obj.to_owned(), field)?
.ok_or_else(|| vm.new_type_error(format!(r#"required field "{field}" missing from {typ}"#)))
Expand Down Expand Up @@ -772,7 +784,7 @@ pub(crate) fn validate_ast_object(vm: &VirtualMachine, object: PyObjectRef) -> P
}

// Used by builtins::compile()
pub const PY_COMPILE_FLAG_AST_ONLY: i32 = 0x0400;
pub const PY_CF_ONLY_AST: i32 = 0x0400;

// The following flags match the values from Include/cpython/compile.h
// Caveat emptor: These flags are undocumented on purpose and depending
Expand All @@ -781,7 +793,7 @@ pub const PY_CF_SOURCE_IS_UTF8: i32 = 0x0100;
pub const PY_CF_DONT_IMPLY_DEDENT: i32 = 0x200;
pub const PY_CF_IGNORE_COOKIE: i32 = 0x0800;
pub const PY_CF_ALLOW_INCOMPLETE_INPUT: i32 = 0x4000;
pub const PY_CF_OPTIMIZED_AST: i32 = 0x8000 | PY_COMPILE_FLAG_AST_ONLY;
pub const PY_CF_OPTIMIZED_AST: i32 = 0x8000 | PY_CF_ONLY_AST;
pub const PY_CF_TYPE_COMMENTS: i32 = 0x1000;
pub const PY_CF_ALLOW_TOP_LEVEL_AWAIT: i32 = 0x2000;

Expand All @@ -802,7 +814,7 @@ const CO_FUTURE_GENERATOR_STOP: i32 = 0x800000;
const CO_FUTURE_ANNOTATIONS: i32 = 0x1000000;

// Used by builtins::compile() - the summary of all flags
pub const PY_COMPILE_FLAGS_MASK: i32 = PY_COMPILE_FLAG_AST_ONLY
pub const PY_COMPILE_FLAGS_MASK: i32 = PY_CF_ONLY_AST
| PY_CF_SOURCE_IS_UTF8
| PY_CF_DONT_IMPLY_DEDENT
| PY_CF_IGNORE_COOKIE
Expand Down
48 changes: 24 additions & 24 deletions crates/vm/src/stdlib/ast/exception.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,22 +9,22 @@ impl Node for ast::ExceptHandler {
}
}
fn ast_from_object(
_vm: &VirtualMachine,
vm: &VirtualMachine,
source_file: &SourceFile,
_object: PyObjectRef,
object: PyObjectRef,
) -> PyResult<Self> {
let _cls = _object.class();
let cls = object.class();
Ok(
if _cls.is(pyast::NodeExceptHandlerExceptHandler::static_type()) {
if cls.is(pyast::NodeExceptHandlerExceptHandler::static_type()) {
Self::ExceptHandler(ast::ExceptHandlerExceptHandler::ast_from_object(
_vm,
vm,
source_file,
_object,
object,
)?)
} else {
return Err(_vm.new_type_error(format!(
return Err(vm.new_type_error(format!(
"expected some sort of excepthandler, but got {}",
_object.repr(_vm)?
object.repr(vm)?
)));
},
)
Expand All @@ -33,50 +33,50 @@ impl Node for ast::ExceptHandler {

// constructor
impl Node for ast::ExceptHandlerExceptHandler {
fn ast_to_object(self, _vm: &VirtualMachine, source_file: &SourceFile) -> PyObjectRef {
fn ast_to_object(self, vm: &VirtualMachine, source_file: &SourceFile) -> PyObjectRef {
let Self {
node_index: _,
type_,
name,
body,
range: _range,
range,
} = self;
let node = NodeAst
.into_ref_with_type(
_vm,
vm,
pyast::NodeExceptHandlerExceptHandler::static_type().to_owned(),
)
.unwrap();
let dict = node.as_object().dict().unwrap();
dict.set_item("type", type_.ast_to_object(_vm, source_file), _vm)
dict.set_item("type", type_.ast_to_object(vm, source_file), vm)
.unwrap();
dict.set_item("name", name.ast_to_object(_vm, source_file), _vm)
dict.set_item("name", name.ast_to_object(vm, source_file), vm)
.unwrap();
dict.set_item("body", body.ast_to_object(_vm, source_file), _vm)
dict.set_item("body", body.ast_to_object(vm, source_file), vm)
.unwrap();
node_add_location(&dict, _range, _vm, source_file);
node_add_location(&dict, range, vm, source_file);
node.into()
}

fn ast_from_object(
_vm: &VirtualMachine,
vm: &VirtualMachine,
source_file: &SourceFile,
_object: PyObjectRef,
object: PyObjectRef,
) -> PyResult<Self> {
Ok(Self {
node_index: Default::default(),
type_: get_node_field_opt(_vm, &_object, "type")?
.map(|obj| Node::ast_from_object(_vm, source_file, obj))
type_: get_node_field_opt(vm, &object, "type")?
.map(|obj| Node::ast_from_object(vm, source_file, obj))
.transpose()?,
name: get_node_field_opt(_vm, &_object, "name")?
.map(|obj| Node::ast_from_object(_vm, source_file, obj))
name: get_node_field_opt(vm, &object, "name")?
.map(|obj| Node::ast_from_object(vm, source_file, obj))
.transpose()?,
body: Node::ast_from_object(
_vm,
vm,
source_file,
get_node_field(_vm, &_object, "body", "ExceptHandler")?,
get_node_field(vm, &object, "body", "ExceptHandler")?,
)?,
range: range_from_object(_vm, source_file, _object, "ExceptHandler")?,
range: range_from_object(vm, source_file, object, "ExceptHandler")?,
})
}
}
16 changes: 5 additions & 11 deletions crates/vm/src/stdlib/ast/expression.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1256,26 +1256,20 @@ impl Node for ast::ExprContext {
unimplemented!("Invalid expression context is not allowed in Python AST")
}
};
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
return instance;
}
NodeAst
.into_ref_with_type(vm, node_type.to_owned())
.unwrap()
.into()
singleton_node_to_object(vm, node_type)
}

fn ast_from_object(
vm: &VirtualMachine,
_source_file: &SourceFile,
object: PyObjectRef,
) -> PyResult<Self> {
let _cls = object.class();
Ok(if _cls.is(pyast::NodeExprContextLoad::static_type()) {
let cls = object.class();
Ok(if cls.is(pyast::NodeExprContextLoad::static_type()) {
Self::Load
} else if _cls.is(pyast::NodeExprContextStore::static_type()) {
} else if cls.is(pyast::NodeExprContextStore::static_type()) {
Self::Store
} else if _cls.is(pyast::NodeExprContextDel::static_type()) {
} else if cls.is(pyast::NodeExprContextDel::static_type()) {
Self::Del
} else {
return Err(vm.new_type_error(format!(
Expand Down
130 changes: 53 additions & 77 deletions crates/vm/src/stdlib/ast/operator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,29 +8,23 @@ impl Node for ast::BoolOp {
Self::And => pyast::NodeBoolOpAnd::static_type(),
Self::Or => pyast::NodeBoolOpOr::static_type(),
};
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
return instance;
}
NodeAst
.into_ref_with_type(vm, node_type.to_owned())
.unwrap()
.into()
singleton_node_to_object(vm, node_type)
}

fn ast_from_object(
_vm: &VirtualMachine,
vm: &VirtualMachine,
_source_file: &SourceFile,
_object: PyObjectRef,
object: PyObjectRef,
) -> PyResult<Self> {
let _cls = _object.class();
Ok(if _cls.is(pyast::NodeBoolOpAnd::static_type()) {
let cls = object.class();
Ok(if cls.is(pyast::NodeBoolOpAnd::static_type()) {
Self::And
} else if _cls.is(pyast::NodeBoolOpOr::static_type()) {
} else if cls.is(pyast::NodeBoolOpOr::static_type()) {
Self::Or
} else {
return Err(_vm.new_type_error(format!(
return Err(vm.new_type_error(format!(
"expected some sort of boolop, but got {}",
_object.repr(_vm)?
object.repr(vm)?
)));
})
}
Expand All @@ -54,51 +48,45 @@ impl Node for ast::Operator {
Self::BitAnd => pyast::NodeOperatorBitAnd::static_type(),
Self::FloorDiv => pyast::NodeOperatorFloorDiv::static_type(),
};
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
return instance;
}
NodeAst
.into_ref_with_type(vm, node_type.to_owned())
.unwrap()
.into()
singleton_node_to_object(vm, node_type)
}

fn ast_from_object(
_vm: &VirtualMachine,
vm: &VirtualMachine,
_source_file: &SourceFile,
_object: PyObjectRef,
object: PyObjectRef,
) -> PyResult<Self> {
let _cls = _object.class();
Ok(if _cls.is(pyast::NodeOperatorAdd::static_type()) {
let cls = object.class();
Ok(if cls.is(pyast::NodeOperatorAdd::static_type()) {
Self::Add
} else if _cls.is(pyast::NodeOperatorSub::static_type()) {
} else if cls.is(pyast::NodeOperatorSub::static_type()) {
Self::Sub
} else if _cls.is(pyast::NodeOperatorMult::static_type()) {
} else if cls.is(pyast::NodeOperatorMult::static_type()) {
Self::Mult
} else if _cls.is(pyast::NodeOperatorMatMult::static_type()) {
} else if cls.is(pyast::NodeOperatorMatMult::static_type()) {
Self::MatMult
} else if _cls.is(pyast::NodeOperatorDiv::static_type()) {
} else if cls.is(pyast::NodeOperatorDiv::static_type()) {
Self::Div
} else if _cls.is(pyast::NodeOperatorMod::static_type()) {
} else if cls.is(pyast::NodeOperatorMod::static_type()) {
Self::Mod
} else if _cls.is(pyast::NodeOperatorPow::static_type()) {
} else if cls.is(pyast::NodeOperatorPow::static_type()) {
Self::Pow
} else if _cls.is(pyast::NodeOperatorLShift::static_type()) {
} else if cls.is(pyast::NodeOperatorLShift::static_type()) {
Self::LShift
} else if _cls.is(pyast::NodeOperatorRShift::static_type()) {
} else if cls.is(pyast::NodeOperatorRShift::static_type()) {
Self::RShift
} else if _cls.is(pyast::NodeOperatorBitOr::static_type()) {
} else if cls.is(pyast::NodeOperatorBitOr::static_type()) {
Self::BitOr
} else if _cls.is(pyast::NodeOperatorBitXor::static_type()) {
} else if cls.is(pyast::NodeOperatorBitXor::static_type()) {
Self::BitXor
} else if _cls.is(pyast::NodeOperatorBitAnd::static_type()) {
} else if cls.is(pyast::NodeOperatorBitAnd::static_type()) {
Self::BitAnd
} else if _cls.is(pyast::NodeOperatorFloorDiv::static_type()) {
} else if cls.is(pyast::NodeOperatorFloorDiv::static_type()) {
Self::FloorDiv
} else {
return Err(_vm.new_type_error(format!(
return Err(vm.new_type_error(format!(
"expected some sort of operator, but got {}",
_object.repr(_vm)?
object.repr(vm)?
)));
})
}
Expand All @@ -113,33 +101,27 @@ impl Node for ast::UnaryOp {
Self::UAdd => pyast::NodeUnaryOpUAdd::static_type(),
Self::USub => pyast::NodeUnaryOpUSub::static_type(),
};
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
return instance;
}
NodeAst
.into_ref_with_type(vm, node_type.to_owned())
.unwrap()
.into()
singleton_node_to_object(vm, node_type)
}

fn ast_from_object(
_vm: &VirtualMachine,
vm: &VirtualMachine,
_source_file: &SourceFile,
_object: PyObjectRef,
object: PyObjectRef,
) -> PyResult<Self> {
let _cls = _object.class();
Ok(if _cls.is(pyast::NodeUnaryOpInvert::static_type()) {
let cls = object.class();
Ok(if cls.is(pyast::NodeUnaryOpInvert::static_type()) {
Self::Invert
} else if _cls.is(pyast::NodeUnaryOpNot::static_type()) {
} else if cls.is(pyast::NodeUnaryOpNot::static_type()) {
Self::Not
} else if _cls.is(pyast::NodeUnaryOpUAdd::static_type()) {
} else if cls.is(pyast::NodeUnaryOpUAdd::static_type()) {
Self::UAdd
} else if _cls.is(pyast::NodeUnaryOpUSub::static_type()) {
} else if cls.is(pyast::NodeUnaryOpUSub::static_type()) {
Self::USub
} else {
return Err(_vm.new_type_error(format!(
return Err(vm.new_type_error(format!(
"expected some sort of unaryop, but got {}",
_object.repr(_vm)?
object.repr(vm)?
)));
})
}
Expand All @@ -160,45 +142,39 @@ impl Node for ast::CmpOp {
Self::In => pyast::NodeCmpOpIn::static_type(),
Self::NotIn => pyast::NodeCmpOpNotIn::static_type(),
};
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
return instance;
}
NodeAst
.into_ref_with_type(vm, node_type.to_owned())
.unwrap()
.into()
singleton_node_to_object(vm, node_type)
}

fn ast_from_object(
_vm: &VirtualMachine,
vm: &VirtualMachine,
_source_file: &SourceFile,
_object: PyObjectRef,
object: PyObjectRef,
) -> PyResult<Self> {
let _cls = _object.class();
Ok(if _cls.is(pyast::NodeCmpOpEq::static_type()) {
let cls = object.class();
Ok(if cls.is(pyast::NodeCmpOpEq::static_type()) {
Self::Eq
} else if _cls.is(pyast::NodeCmpOpNotEq::static_type()) {
} else if cls.is(pyast::NodeCmpOpNotEq::static_type()) {
Self::NotEq
} else if _cls.is(pyast::NodeCmpOpLt::static_type()) {
} else if cls.is(pyast::NodeCmpOpLt::static_type()) {
Self::Lt
} else if _cls.is(pyast::NodeCmpOpLtE::static_type()) {
} else if cls.is(pyast::NodeCmpOpLtE::static_type()) {
Self::LtE
} else if _cls.is(pyast::NodeCmpOpGt::static_type()) {
} else if cls.is(pyast::NodeCmpOpGt::static_type()) {
Self::Gt
} else if _cls.is(pyast::NodeCmpOpGtE::static_type()) {
} else if cls.is(pyast::NodeCmpOpGtE::static_type()) {
Self::GtE
} else if _cls.is(pyast::NodeCmpOpIs::static_type()) {
} else if cls.is(pyast::NodeCmpOpIs::static_type()) {
Self::Is
} else if _cls.is(pyast::NodeCmpOpIsNot::static_type()) {
} else if cls.is(pyast::NodeCmpOpIsNot::static_type()) {
Self::IsNot
} else if _cls.is(pyast::NodeCmpOpIn::static_type()) {
} else if cls.is(pyast::NodeCmpOpIn::static_type()) {
Self::In
} else if _cls.is(pyast::NodeCmpOpNotIn::static_type()) {
} else if cls.is(pyast::NodeCmpOpNotIn::static_type()) {
Self::NotIn
} else {
return Err(_vm.new_type_error(format!(
return Err(vm.new_type_error(format!(
"expected some sort of cmpop, but got {}",
_object.repr(_vm)?
object.repr(vm)?
)));
})
}
Expand Down
Loading
Loading