Refactor builtin_print to avoid duplicate loop code
@@ -563,56 +563,73 @@ pub struct PrintOptions {
file: Option<PyObjectRef>,
}
-pub fn builtin_print(objects: Args, options: PrintOptions, vm: &VirtualMachine) -> PyResult<()> {
- if let Some(file) = options.file {
- let sep = match options.sep {
- Some(sep) => sep.into_pyobject(vm).unwrap(),
- _ => PyString::from(" ").into_pyobject(vm).unwrap(),
- };
+trait Printer {
+ fn write(&mut self, vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<()>;
+ fn flush(&mut self, vm: &VirtualMachine) -> PyResult<()>;
+}
- let mut first = true;
- for object in objects {
- if first {
- first = false;
- } else {
- vm.call_method(&file, "write", vec![sep.clone()])?;
+impl Printer for &'_ PyObjectRef {
+ fn write(&mut self, vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<()> {
+ vm.call_method(self, "write", vec![obj])?;
+ Ok(())
}
- vm.call_method(&file, "write", vec![object])?;
+ fn flush(&mut self, vm: &VirtualMachine) -> PyResult<()> {
+ vm.call_method(self, "flush", vec![])?;
+ Ok(())
+ }
}
- let end = match options.end {
- Some(end) => end.into_pyobject(vm).unwrap(),
- _ => PyString::from("\n").into_pyobject(vm).unwrap(),
- };
- vm.call_method(&file, "write", vec![end])?;
+impl Printer for std::io::StdoutLock<'_> {
+ fn write(&mut self, vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<()> {
+ let s = &vm.to_str(&obj)?.value;
+ write!(self, "{}", s).unwrap();
+ Ok(())
+ }
- if options.flush {
- vm.call_method(&file, "flush", vec![])?;
+ fn flush(&mut self, _vm: &VirtualMachine) -> PyResult<()> {
+ <Self as std::io::Write>::flush(self).unwrap();
+ Ok(())
+ }
}
- } else {
- let sep = options.sep.as_ref().map_or(" ", |sep| &sep.value);
+pub fn builtin_print(objects: Args, options: PrintOptions, vm: &VirtualMachine) -> PyResult<()> {
let stdout = io::stdout();
- let mut stdout_lock = stdout.lock();
+
+ let mut printer: Box<dyn Printer> = if let Some(file) = &options.file {
+ Box::new(file)
+ } else {
+ Box::new(stdout.lock())
+ };
+
+ let sep = options
+ .sep
+ .as_ref()
+ .map_or(" ", |sep| &sep.value)
+ .into_pyobject(vm)
+ .unwrap();
+
let mut first = true;
for object in objects {
if first {
first = false;
} else {
- write!(stdout_lock, "{}", sep).unwrap();
+ printer.write(vm, sep.clone())?;
}
- let s = &vm.to_str(&object)?.value;
- write!(stdout_lock, "{}", s).unwrap();
+ printer.write(vm, object)?;
}
- let end = options.end.as_ref().map_or("\n", |end| &end.value);
- write!(stdout_lock, "{}", end).unwrap();
+ let end = options
+ .end
+ .as_ref()
+ .map_or("\n", |end| &end.value)
+ .into_pyobject(vm)
+ .unwrap();
+ printer.write(vm, end)?;
if options.flush {
- stdout_lock.flush().unwrap();
- }
+ printer.flush(vm)?;
}
Ok(())
