diff --git a/stdlib/src/dis.rs b/stdlib/src/dis.rs index 341137f91f..67d6c41cb4 100644 --- a/stdlib/src/dis.rs +++ b/stdlib/src/dis.rs @@ -2,14 +2,24 @@ pub(crate) use decl::make_module; #[pymodule(name = "dis")] mod decl { - use crate::vm::{ + use rustpython_vm::{ PyObjectRef, PyRef, PyResult, TryFromObject, VirtualMachine, builtins::{PyCode, PyDictRef, PyStrRef}, bytecode::CodeFlags, + function::OptionalArg, }; + #[derive(FromArgs)] + struct DisArgs { + #[pyarg(positional)] + obj: PyObjectRef, + #[pyarg(any, optional)] + file: OptionalArg, + } + #[pyfunction] - fn dis(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + fn dis(args: DisArgs, vm: &VirtualMachine) -> PyResult<()> { + let DisArgs { obj, file } = args; let co = if let Ok(co) = obj.get_attr("__code__", vm) { // Method or function: PyRef::try_from_object(vm, co)? @@ -33,12 +43,41 @@ mod decl { } else { PyRef::try_from_object(vm, obj)? }; - disassemble(co) + disassemble_to_file(co, file.into_option(), vm) } #[pyfunction] - fn disassemble(co: PyRef) -> PyResult<()> { - print!("{}", &co.code); + fn disassemble( + co: PyRef, + file: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult<()> { + disassemble_to_file(co, file.into_option(), vm) + } + + fn disassemble_to_file( + co: PyRef, + file: Option, + vm: &VirtualMachine, + ) -> PyResult<()> { + let output = format!("{}", &co.code); + + match file { + Some(file_obj) => { + // Write to the provided file object + if let Ok(write_method) = file_obj.get_attr("write", vm) { + write_method.call((output,), vm)?; + } else { + return Err( + vm.new_type_error("file argument must have a write method".to_owned()) + ); + } + } + None => { + // Write to stdout + print!("{output}"); + } + } Ok(()) }