diff --git a/vm/src/stdlib/dis.rs b/vm/src/stdlib/dis.rs index 2f591f855..4436d1093 100644 --- a/vm/src/stdlib/dis.rs +++ b/vm/src/stdlib/dis.rs @@ -1,57 +1,56 @@ -use crate::bytecode::CodeFlags; -use crate::obj::objcode::PyCodeRef; -use crate::obj::objstr::PyStringRef; -use crate::pyobject::{BorrowValue, ItemProtocol, PyObjectRef, PyResult, TryFromObject}; -use crate::vm::VirtualMachine; -use rustpython_compiler::compile; +pub(crate) use decl::make_module; -fn dis_dis(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult { - // Method or function: - if let Ok(co) = vm.get_attribute(obj.clone(), "__code__") { - return dis_disassemble(co, vm); +#[pymodule(name = "dis")] +mod decl { + use crate::bytecode::CodeFlags; + use crate::obj::objcode::PyCodeRef; + use crate::obj::objdict::PyDictRef; + use crate::obj::objstr::PyStringRef; + use crate::pyobject::{BorrowValue, ItemProtocol, PyObjectRef, PyResult, TryFromObject}; + use crate::vm::VirtualMachine; + use rustpython_compiler::compile; + + #[pyfunction] + fn dis(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult { + // Method or function: + if let Ok(co) = vm.get_attribute(obj.clone(), "__code__") { + return disassemble(co, vm); + } + + // String: + if let Ok(co_str) = PyStringRef::try_from_object(vm, obj.clone()) { + let code = vm + .compile( + co_str.borrow_value(), + compile::Mode::Exec, + "".to_owned(), + ) + .map_err(|err| vm.new_syntax_error(&err))? + .into_object(); + return disassemble(code, vm); + } + + disassemble(obj, vm) } - // String: - if let Ok(co_str) = PyStringRef::try_from_object(vm, obj.clone()) { - let code = vm - .compile( - co_str.borrow_value(), - compile::Mode::Exec, - "".to_owned(), + #[pyfunction] + fn disassemble(co: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let code = &PyCodeRef::try_from_object(vm, co)?.code; + print!("{}", code); + Ok(vm.get_none()) + } + + #[pyattr(name = "COMPILER_FLAG_NAMES")] + fn compiler_flag_names(vm: &VirtualMachine) -> PyDictRef { + let dict = vm.ctx.new_dict(); + for (name, flag) in CodeFlags::NAME_MAPPING { + dict.set_item( + vm.ctx.new_int(flag.bits()), + vm.ctx.new_str((*name).to_owned()), + vm, ) - .map_err(|err| vm.new_syntax_error(&err))? - .into_object(); - return dis_disassemble(code, vm); + .unwrap(); + } + dict } - - dis_disassemble(obj, vm) -} - -fn dis_disassemble(co: PyObjectRef, vm: &VirtualMachine) -> PyResult { - let code = &PyCodeRef::try_from_object(vm, co)?.code; - print!("{}", code); - Ok(vm.get_none()) -} - -fn dis_compiler_flag_names(vm: &VirtualMachine) -> PyObjectRef { - let dict = vm.ctx.new_dict(); - for (name, flag) in CodeFlags::NAME_MAPPING { - dict.set_item( - vm.ctx.new_int(flag.bits()), - vm.ctx.new_str((*name).to_owned()), - vm, - ) - .unwrap(); - } - dict.into_object() -} - -pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { - let ctx = &vm.ctx; - - py_module!(vm, "dis", { - "dis" => ctx.new_function(dis_dis), - "disassemble" => ctx.new_function(dis_disassemble), - "COMPILER_FLAG_NAMES" => dis_compiler_flag_names(vm), - }) }