rustpython_vm/
py_io.rs

1use core::{fmt, ops};
2use std::io;
3
4use crate::{
5    PyObject, PyObjectRef, PyResult, VirtualMachine,
6    builtins::{PyBaseExceptionRef, PyBytes, PyStr},
7    common::ascii,
8};
9
10pub trait Write {
11    type Error;
12
13    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> Result<(), Self::Error>;
14}
15
16#[repr(transparent)]
17pub struct IoWriter<T>(pub T);
18
19impl<T> IoWriter<T> {
20    pub fn from_ref(x: &mut T) -> &mut Self {
21        // SAFETY: IoWriter is repr(transparent) over T
22        unsafe { &mut *(x as *mut T as *mut Self) }
23    }
24}
25
26impl<T> ops::Deref for IoWriter<T> {
27    type Target = T;
28
29    fn deref(&self) -> &T {
30        &self.0
31    }
32}
33
34impl<T> ops::DerefMut for IoWriter<T> {
35    fn deref_mut(&mut self) -> &mut T {
36        &mut self.0
37    }
38}
39
40impl<W> Write for IoWriter<W>
41where
42    W: io::Write,
43{
44    type Error = io::Error;
45
46    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> io::Result<()> {
47        <W as io::Write>::write_fmt(&mut self.0, args)
48    }
49}
50
51impl Write for String {
52    type Error = fmt::Error;
53
54    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result {
55        <Self as fmt::Write>::write_fmt(self, args)
56    }
57}
58
59pub struct PyWriter<'vm>(pub PyObjectRef, pub &'vm VirtualMachine);
60
61impl Write for PyWriter<'_> {
62    type Error = PyBaseExceptionRef;
63
64    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> Result<(), Self::Error> {
65        let Self(obj, vm) = self;
66
67        vm.call_method(obj, "write", (args.to_string(),)).map(drop)
68    }
69}
70
71pub fn file_readline(obj: &PyObject, size: Option<usize>, vm: &VirtualMachine) -> PyResult {
72    let args = size.map_or_else(Vec::new, |size| vec![vm.ctx.new_int(size).into()]);
73    let ret = vm.call_method(obj, "readline", args)?;
74    let eof_err = || {
75        vm.new_exception(
76            vm.ctx.exceptions.eof_error.to_owned(),
77            vec![vm.ctx.new_str(ascii!("EOF when reading a line")).into()],
78        )
79    };
80
81    let ret = match_class!(match ret {
82        s @ PyStr => {
83            // Use as_wtf8() to handle strings with surrogates (e.g., surrogateescape)
84            let s_wtf8 = s.as_wtf8();
85            if s_wtf8.is_empty() {
86                return Err(eof_err());
87            }
88
89            // '\n' is ASCII, so we can check bytes directly
90            if s_wtf8.as_bytes().last() == Some(&b'\n') {
91                let no_nl = &s_wtf8[..s_wtf8.len() - 1];
92                vm.ctx.new_str(no_nl).into()
93            } else {
94                s.into()
95            }
96        }
97        b @ PyBytes => {
98            let buf = b.as_bytes();
99            if buf.is_empty() {
100                return Err(eof_err());
101            }
102
103            if buf.last() == Some(&b'\n') {
104                vm.ctx.new_bytes(buf[..buf.len() - 1].to_owned()).into()
105            } else {
106                b.into()
107            }
108        }
109        _ => return Err(vm.new_type_error("object.readline() returned non-string")),
110    });
111    Ok(ret)
112}
rustpython_vm/py_io.rs

rustpython_vm/
py_io.rs
