1pub(crate) use _io::module_def;
5#[cfg(all(unix, feature = "threading"))]
6pub(crate) use _io::reinit_std_streams_after_fork;
7
8cfg_if::cfg_if! {
9 if #[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))] {
10 use crate::common::crt_fd::Offset;
11 } else {
12 type Offset = i64;
13 }
14}
15
16cfg_if::cfg_if! {
18 if #[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))] {
19 const EAGAIN: i32 = libc::EAGAIN;
20 } else {
21 const EAGAIN: i32 = 11; }
23}
24
25use crate::{
26 AsObject, PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine, builtins::PyModule,
27};
28pub use _io::{OpenArgs, io_open as open};
29
30fn file_closed(file: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
31 file.get_attr("closed", vm)?.try_to_bool(vm)
32}
33
34const DEFAULT_BUFFER_SIZE: usize = 128 * 1024;
35
36fn iobase_finalize(zelf: &PyObject, vm: &VirtualMachine) {
38 let closed = match vm.get_attribute_opt(zelf.to_owned(), "closed") {
41 Ok(Some(val)) => match val.try_to_bool(vm) {
42 Ok(b) => b,
43 Err(_) => return,
44 },
45 _ => return,
46 };
47 if !closed {
48 let _ = zelf.set_attr("_finalizing", vm.ctx.true_value.clone(), vm);
51 if let Err(e) = vm.call_method(zelf, "close", ()) {
52 if !e.fast_isinstance(vm.ctx.exceptions.broken_pipe_error) {
58 vm.run_unraisable(e, None, zelf.to_owned());
59 }
60 }
61 }
62}
63
64#[derive(Copy, Clone)]
66#[repr(transparent)]
67pub struct Fildes(pub i32);
68
69impl TryFromObject for Fildes {
70 fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
71 use crate::builtins::int;
72 let int = match obj.downcast::<int::PyInt>() {
73 Ok(i) => i,
74 Err(obj) => {
75 let fileno_meth = vm.get_attribute_opt(obj, "fileno")?.ok_or_else(|| {
76 vm.new_type_error("argument must be an int, or have a fileno() method.")
77 })?;
78 fileno_meth
79 .call((), vm)?
80 .downcast()
81 .map_err(|_| vm.new_type_error("fileno() returned a non-integer"))?
82 }
83 };
84 let fd = int.try_to_primitive(vm)?;
85 if fd < 0 {
86 return Err(vm.new_value_error(format!(
87 "file descriptor cannot be a negative integer ({fd})"
88 )));
89 }
90 Ok(Self(fd))
91 }
92}
93
94#[cfg(unix)]
95impl std::os::fd::AsFd for Fildes {
96 fn as_fd(&self) -> std::os::fd::BorrowedFd<'_> {
97 unsafe { std::os::fd::BorrowedFd::borrow_raw(self.0) }
100 }
101}
102#[cfg(unix)]
103impl std::os::fd::AsRawFd for Fildes {
104 fn as_raw_fd(&self) -> std::os::fd::RawFd {
105 self.0
106 }
107}
108
109#[pymodule]
110mod _io {
111 use super::*;
112 use crate::{
113 AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
114 TryFromBorrowedObject, TryFromObject,
115 builtins::{
116 PyBaseExceptionRef, PyBool, PyByteArray, PyBytes, PyBytesRef, PyDict, PyMemoryView,
117 PyStr, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef, PyUtf8Str, PyUtf8StrRef,
118 },
119 class::StaticType,
120 common::lock::{
121 PyMappedThreadMutexGuard, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard,
122 PyThreadMutex, PyThreadMutexGuard,
123 },
124 common::wtf8::{Wtf8, Wtf8Buf},
125 convert::ToPyObject,
126 exceptions::cstring_error,
127 function::{
128 ArgBytesLike, ArgIterable, ArgMemoryBuffer, ArgSize, Either, FuncArgs, IntoFuncArgs,
129 OptionalArg, OptionalOption, PySetterValue,
130 },
131 protocol::{
132 BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn, VecBuffer,
133 },
134 recursion::ReprGuard,
135 types::{
136 Callable, Constructor, DefaultConstructor, Destructor, Initializer, IterNext, Iterable,
137 Representable,
138 },
139 vm::VirtualMachine,
140 };
141 use alloc::borrow::Cow;
142 use bstr::ByteSlice;
143 use core::{
144 ops::Range,
145 sync::atomic::{AtomicBool, Ordering},
146 };
147 use crossbeam_utils::atomic::AtomicCell;
148 use malachite_bigint::BigInt;
149 use num_traits::ToPrimitive;
150 use std::io::{self, Cursor, SeekFrom, prelude::*};
151
152 #[allow(clippy::let_and_return)]
153 fn validate_whence(whence: i32) -> bool {
154 let x = (0..=2).contains(&whence);
155 cfg_if::cfg_if! {
156 if #[cfg(any(target_os = "dragonfly", target_os = "freebsd", target_os = "linux"))] {
157 x || matches!(whence, libc::SEEK_DATA | libc::SEEK_HOLE)
158 } else {
159 x
160 }
161 }
162 }
163
164 fn ensure_unclosed(file: &PyObject, msg: &str, vm: &VirtualMachine) -> PyResult<()> {
165 if file.get_attr("closed", vm)?.try_to_bool(vm)? {
166 Err(vm.new_value_error(msg))
167 } else {
168 Ok(())
169 }
170 }
171
172 #[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))]
177 fn trap_eintr<T>(result: PyResult<T>, vm: &VirtualMachine) -> PyResult<Option<T>> {
178 match result {
179 Ok(val) => Ok(Some(val)),
180 Err(exc) => {
181 if exc.fast_isinstance(vm.ctx.exceptions.os_error)
183 && let Ok(errno_attr) = exc.as_object().get_attr("errno", vm)
184 && let Ok(errno_val) = i32::try_from_object(vm, errno_attr)
185 && errno_val == libc::EINTR
186 {
187 vm.check_signals()?;
188 return Ok(None);
189 }
190 Err(exc)
191 }
192 }
193 }
194
195 #[cfg(all(target_arch = "wasm32", not(target_os = "wasi")))]
197 fn trap_eintr<T>(result: PyResult<T>, _vm: &VirtualMachine) -> PyResult<Option<T>> {
198 result.map(Some)
199 }
200
201 pub fn new_unsupported_operation(vm: &VirtualMachine, msg: String) -> PyBaseExceptionRef {
202 vm.new_os_subtype_error(unsupported_operation().to_owned(), None, msg)
203 .upcast()
204 }
205
206 fn _unsupported<T>(vm: &VirtualMachine, zelf: &PyObject, operation: &str) -> PyResult<T> {
207 Err(new_unsupported_operation(
208 vm,
209 format!("{}.{}() not supported", zelf.class().name(), operation),
210 ))
211 }
212
213 #[derive(FromArgs)]
214 pub(super) struct OptionalSize {
215 #[pyarg(positional, default)]
218 size: Option<ArgSize>,
219 }
220
221 impl OptionalSize {
222 #[allow(clippy::wrong_self_convention)]
223 pub fn to_usize(self) -> Option<usize> {
224 self.size?.to_usize()
225 }
226
227 pub fn try_usize(self, vm: &VirtualMachine) -> PyResult<Option<usize>> {
228 self.size
229 .map(|v| {
230 let v = *v;
231 if v >= 0 {
232 Ok(v as usize)
233 } else {
234 Err(vm.new_value_error(format!("Negative size value {v}")))
235 }
236 })
237 .transpose()
238 }
239 }
240
241 fn os_err(vm: &VirtualMachine, err: io::Error) -> PyBaseExceptionRef {
242 use crate::convert::ToPyException;
243 err.to_pyexception(vm)
244 }
245
246 pub(super) fn io_closed_error(vm: &VirtualMachine) -> PyBaseExceptionRef {
247 vm.new_value_error("I/O operation on closed file")
248 }
249
250 #[pyattr]
251 const DEFAULT_BUFFER_SIZE: usize = super::DEFAULT_BUFFER_SIZE;
252
253 pub(super) fn seekfrom(
254 vm: &VirtualMachine,
255 offset: PyObjectRef,
256 how: OptionalArg<i32>,
257 ) -> PyResult<SeekFrom> {
258 let seek = match how {
259 OptionalArg::Present(0) | OptionalArg::Missing => {
260 SeekFrom::Start(offset.try_into_value(vm)?)
261 }
262 OptionalArg::Present(1) => SeekFrom::Current(offset.try_into_value(vm)?),
263 OptionalArg::Present(2) => SeekFrom::End(offset.try_into_value(vm)?),
264 _ => return Err(vm.new_value_error("invalid value for how")),
265 };
266 Ok(seek)
267 }
268
269 #[derive(Debug)]
270 struct BufferedIO {
271 cursor: Cursor<Vec<u8>>,
272 }
273
274 impl BufferedIO {
275 const fn new(cursor: Cursor<Vec<u8>>) -> Self {
276 Self { cursor }
277 }
278
279 fn write(&mut self, data: &[u8]) -> Option<u64> {
280 if data.is_empty() {
281 return Some(0);
282 }
283 let length = data.len();
284 self.cursor.write_all(data).ok()?;
285 Some(length as u64)
286 }
287
288 fn getvalue(&self) -> Vec<u8> {
290 self.cursor.clone().into_inner()
291 }
292
293 fn seek(&mut self, seek: SeekFrom) -> io::Result<u64> {
295 self.cursor.seek(seek)
296 }
297
298 fn read(&mut self, bytes: Option<usize>) -> Option<Vec<u8>> {
300 let pos = self.cursor.position().to_usize()?;
301 let avail_slice = self.cursor.get_ref().get(pos..)?;
302 let n = bytes.map_or_else(
304 || avail_slice.len(),
305 |n| core::cmp::min(n, avail_slice.len()),
306 );
307 let b = avail_slice[..n].to_vec();
308 self.cursor.set_position((pos + n) as u64);
309 Some(b)
310 }
311
312 const fn tell(&self) -> u64 {
313 self.cursor.position()
314 }
315
316 fn readline(&mut self, size: Option<usize>, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
317 self.read_until(size, b'\n', vm)
318 }
319
320 fn read_until(
321 &mut self,
322 size: Option<usize>,
323 byte: u8,
324 vm: &VirtualMachine,
325 ) -> PyResult<Vec<u8>> {
326 let size = match size {
327 None => {
328 let mut buf: Vec<u8> = Vec::new();
329 self.cursor
330 .read_until(byte, &mut buf)
331 .map_err(|err| os_err(vm, err))?;
332 return Ok(buf);
333 }
334 Some(0) => {
335 return Ok(Vec::new());
336 }
337 Some(size) => size,
338 };
339
340 let available = {
341 let buf = self.cursor.fill_buf().map_err(|err| os_err(vm, err))?;
344 if size < buf.len() { &buf[..size] } else { buf }
345 };
346 let buf = match available.find_byte(byte) {
347 Some(i) => available[..=i].to_vec(),
348 _ => available.to_vec(),
349 };
350 self.cursor.consume(buf.len());
351 Ok(buf)
352 }
353
354 fn truncate(&mut self, pos: Option<usize>) -> usize {
355 let pos = pos.unwrap_or_else(|| self.tell() as usize);
356 self.cursor.get_mut().truncate(pos);
357 pos
358 }
359 }
360
361 fn check_closed(file: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
362 if file_closed(file, vm)? {
363 Err(io_closed_error(vm))
364 } else {
365 Ok(())
366 }
367 }
368
369 fn check_readable(file: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
370 if vm.call_method(file, "readable", ())?.try_to_bool(vm)? {
371 Ok(())
372 } else {
373 Err(new_unsupported_operation(
374 vm,
375 "File or stream is not readable".to_owned(),
376 ))
377 }
378 }
379
380 fn check_writable(file: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
381 if vm.call_method(file, "writable", ())?.try_to_bool(vm)? {
382 Ok(())
383 } else {
384 Err(new_unsupported_operation(
385 vm,
386 "File or stream is not writable.".to_owned(),
387 ))
388 }
389 }
390
391 fn check_seekable(file: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
392 if vm.call_method(file, "seekable", ())?.try_to_bool(vm)? {
393 Ok(())
394 } else {
395 Err(new_unsupported_operation(
396 vm,
397 "File or stream is not seekable".to_owned(),
398 ))
399 }
400 }
401
402 fn check_decoded(decoded: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyStrRef> {
403 decoded.downcast().map_err(|obj| {
404 vm.new_type_error(format!(
405 "decoder should return a string result, not '{}'",
406 obj.class().name()
407 ))
408 })
409 }
410
411 #[pyattr]
412 #[pyclass(name = "_IOBase")]
413 #[derive(Debug, Default, PyPayload)]
414 pub struct _IOBase;
415
416 #[pyclass(
417 with(IterNext, Iterable, Destructor),
418 flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
419 )]
420 impl _IOBase {
421 #[pymethod]
422 fn seek(
423 zelf: PyObjectRef,
424 _pos: PyObjectRef,
425 _whence: OptionalArg,
426 vm: &VirtualMachine,
427 ) -> PyResult {
428 _unsupported(vm, &zelf, "seek")
429 }
430
431 #[pymethod]
432 fn tell(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
433 vm.call_method(&zelf, "seek", (0, 1))
434 }
435
436 #[pymethod]
437 fn truncate(zelf: PyObjectRef, _pos: OptionalArg, vm: &VirtualMachine) -> PyResult {
438 _unsupported(vm, &zelf, "truncate")
439 }
440
441 #[pymethod]
442 fn fileno(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
443 _unsupported(vm, &zelf, "fileno")
444 }
445
446 #[pyattr]
447 fn __closed(ctx: &Context) -> PyRef<PyBool> {
448 ctx.new_bool(false)
449 }
450
451 #[pymethod]
452 fn __enter__(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult {
453 check_closed(&instance, vm)?;
454 Ok(instance)
455 }
456
457 #[pymethod]
458 fn __exit__(instance: PyObjectRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
459 vm.call_method(&instance, "close", ())?;
460 Ok(())
461 }
462
463 #[pymethod]
464 fn flush(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
465 check_closed(&instance, vm)
467 }
468
469 #[pymethod]
470 fn seekable(_self: PyObjectRef) -> bool {
471 false
472 }
473
474 #[pymethod]
475 fn readable(_self: PyObjectRef) -> bool {
476 false
477 }
478
479 #[pymethod]
480 fn writable(_self: PyObjectRef) -> bool {
481 false
482 }
483
484 #[pymethod]
485 fn isatty(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
486 check_closed(&instance, vm)?;
487 Ok(false)
488 }
489
490 #[pygetset]
491 fn closed(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult {
492 instance.get_attr("__closed", vm)
493 }
494
495 #[pymethod]
496 fn close(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
497 iobase_close(&instance, vm)
498 }
499
500 #[pymethod]
501 fn readline(
502 instance: PyObjectRef,
503 size: OptionalSize,
504 vm: &VirtualMachine,
505 ) -> PyResult<Vec<u8>> {
506 let size = size.to_usize();
507 let read = instance.get_attr("read", vm)?;
508 let mut res = Vec::new();
509 while size.is_none_or(|s| res.len() < s) {
510 let read_res = ArgBytesLike::try_from_object(vm, read.call((1,), vm)?)?;
511 if read_res.with_ref(|b| b.is_empty()) {
512 break;
513 }
514 read_res.with_ref(|b| res.extend_from_slice(b));
515 if res.ends_with(b"\n") {
516 break;
517 }
518 }
519 Ok(res)
520 }
521
522 #[pymethod]
523 fn readlines(
524 instance: PyObjectRef,
525 hint: OptionalOption<isize>,
526 vm: &VirtualMachine,
527 ) -> PyResult<Vec<PyObjectRef>> {
528 let hint = hint.flatten().unwrap_or(-1);
529 if hint <= 0 {
530 return instance.try_to_value(vm);
531 }
532 let hint = hint as usize;
533 let mut ret = Vec::new();
534 let it = ArgIterable::<PyObjectRef>::try_from_object(vm, instance)?;
535 let mut full_len = 0;
536 for line in it.iter(vm)? {
537 let line = line?;
538 let line_len = line.length(vm)?;
539 ret.push(line.clone());
540 full_len += line_len;
541 if full_len > hint {
542 break;
543 }
544 }
545 Ok(ret)
546 }
547
548 #[pymethod]
549 fn writelines(
550 instance: PyObjectRef,
551 lines: ArgIterable,
552 vm: &VirtualMachine,
553 ) -> PyResult<()> {
554 check_closed(&instance, vm)?;
555 for line in lines.iter(vm)? {
556 vm.call_method(&instance, "write", (line?,))?;
557 }
558 Ok(())
559 }
560
561 #[pymethod(name = "_checkClosed")]
562 fn check_closed(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
563 check_closed(&instance, vm)
564 }
565
566 #[pymethod(name = "_checkReadable")]
567 fn check_readable(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
568 check_readable(&instance, vm)
569 }
570
571 #[pymethod(name = "_checkWritable")]
572 fn check_writable(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
573 check_writable(&instance, vm)
574 }
575
576 #[pymethod(name = "_checkSeekable")]
577 fn check_seekable(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
578 check_seekable(&instance, vm)
579 }
580 }
581
582 impl Destructor for _IOBase {
583 fn slot_del(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
584 let _ = vm.call_method(zelf, "close", ());
590 Ok(())
591 }
592
593 #[cold]
594 fn del(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<()> {
595 unreachable!("slot_del is implemented")
596 }
597 }
598
599 impl Iterable for _IOBase {
600 fn slot_iter(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
601 check_closed(&zelf, vm)?;
602 Ok(zelf)
603 }
604
605 fn iter(_zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyResult {
606 unreachable!("slot_iter is implemented")
607 }
608 }
609
610 impl IterNext for _IOBase {
611 fn slot_iternext(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
612 let line = vm.call_method(zelf, "readline", ())?;
613 Ok(if !line.clone().try_to_bool(vm)? {
614 PyIterReturn::StopIteration(None)
615 } else {
616 PyIterReturn::Return(line)
617 })
618 }
619
620 fn next(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<PyIterReturn> {
621 unreachable!("slot_iternext is implemented")
622 }
623 }
624
625 pub(super) fn iobase_close(file: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
626 if !file_closed(file, vm)? {
627 let res = vm.call_method(file, "flush", ());
628 file.set_attr("__closed", vm.new_pyobj(true), vm)?;
629 res?;
630 }
631 Ok(())
632 }
633
634 #[pyattr]
635 #[pyclass(name = "_RawIOBase", base = _IOBase)]
636 #[derive(Debug, Default)]
637 #[repr(transparent)]
638 pub(super) struct _RawIOBase(_IOBase);
639
640 #[pyclass(flags(BASETYPE, HAS_DICT, HAS_WEAKREF))]
641 impl _RawIOBase {
642 #[pymethod]
643 fn read(instance: PyObjectRef, size: OptionalSize, vm: &VirtualMachine) -> PyResult {
644 if let Some(size) = size.to_usize() {
645 let b = PyByteArray::from(vec![0; size]).into_ref(&vm.ctx);
647 let n = <Option<isize>>::try_from_object(
648 vm,
649 vm.call_method(&instance, "readinto", (b.clone(),))?,
650 )?;
651 Ok(match n {
652 None => vm.ctx.none(),
653 Some(n) => {
654 if n < 0 || (n as usize) > size {
656 return Err(vm.new_value_error(format!(
657 "readinto returned {n} outside buffer size {size}"
658 )));
659 }
660 let n = n as usize;
661 let mut bytes = b.borrow_buf_mut();
662 bytes.truncate(n);
663 bytes.clone().to_pyobject(vm)
665 }
666 })
667 } else {
668 vm.call_method(&instance, "readall", ())
669 }
670 }
671
672 #[pymethod]
673 fn readall(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult<Option<Vec<u8>>> {
674 let mut chunks = Vec::new();
675 let mut total_len = 0;
676 loop {
677 let data = loop {
679 let res = vm.call_method(&instance, "read", (DEFAULT_BUFFER_SIZE,));
680 match trap_eintr(res, vm)? {
681 Some(val) => break val,
682 None => continue,
683 }
684 };
685 let data = <Option<PyBytesRef>>::try_from_object(vm, data)?;
686 match data {
687 None => {
688 if chunks.is_empty() {
689 return Ok(None);
690 }
691 break;
692 }
693 Some(b) => {
694 if b.as_bytes().is_empty() {
695 break;
696 }
697 total_len += b.as_bytes().len();
698 chunks.push(b)
699 }
700 }
701 }
702 let mut ret = Vec::with_capacity(total_len);
703 for b in chunks {
704 ret.extend_from_slice(b.as_bytes())
705 }
706 Ok(Some(ret))
707 }
708
709 #[pymethod]
710 fn readinto(_instance: PyObjectRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult {
711 Err(vm.new_not_implemented_error(String::new()))
712 }
713
714 #[pymethod]
715 fn write(_instance: PyObjectRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult {
716 Err(vm.new_not_implemented_error(String::new()))
717 }
718 }
719
720 #[pyattr]
721 #[pyclass(name = "_BufferedIOBase", base = _IOBase)]
722 #[derive(Debug, Default)]
723 #[repr(transparent)]
724 struct _BufferedIOBase(_IOBase);
725
726 #[pyclass(flags(BASETYPE, HAS_WEAKREF))]
727 impl _BufferedIOBase {
728 #[pymethod]
729 fn read(zelf: PyObjectRef, _size: OptionalArg, vm: &VirtualMachine) -> PyResult {
730 _unsupported(vm, &zelf, "read")
731 }
732
733 #[pymethod]
734 fn read1(zelf: PyObjectRef, _size: OptionalArg, vm: &VirtualMachine) -> PyResult {
735 _unsupported(vm, &zelf, "read1")
736 }
737
738 fn _readinto(
739 zelf: PyObjectRef,
740 buf_obj: PyObjectRef,
741 method: &str,
742 vm: &VirtualMachine,
743 ) -> PyResult<usize> {
744 let b = ArgMemoryBuffer::try_from_borrowed_object(vm, &buf_obj)?;
745 let l = b.len();
746 let data = vm.call_method(&zelf, method, (l,))?;
747 if data.is(&buf_obj) {
748 return Ok(l);
749 }
750 let mut buf = b.borrow_buf_mut();
751 let data = ArgBytesLike::try_from_object(vm, data)?;
752 let data = data.borrow_buf();
753 match buf.get_mut(..data.len()) {
754 Some(slice) => {
755 slice.copy_from_slice(&data);
756 Ok(data.len())
757 }
758 None => {
759 Err(vm.new_value_error("readinto: buffer and read data have different lengths"))
760 }
761 }
762 }
763 #[pymethod]
764 fn readinto(zelf: PyObjectRef, b: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
765 Self::_readinto(zelf, b, "read", vm)
766 }
767
768 #[pymethod]
769 fn readinto1(zelf: PyObjectRef, b: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
770 Self::_readinto(zelf, b, "read1", vm)
771 }
772
773 #[pymethod]
774 fn write(zelf: PyObjectRef, _b: PyObjectRef, vm: &VirtualMachine) -> PyResult {
775 _unsupported(vm, &zelf, "write")
776 }
777
778 #[pymethod]
779 fn detach(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
780 _unsupported(vm, &zelf, "detach")
781 }
782 }
783
784 #[pyattr]
786 #[pyclass(name = "_TextIOBase", base = _IOBase)]
787 #[derive(Debug, Default)]
788 #[repr(transparent)]
789 struct _TextIOBase(_IOBase);
790
791 #[pyclass(flags(BASETYPE, HAS_WEAKREF))]
792 impl _TextIOBase {
793 #[pygetset]
794 fn encoding(_zelf: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
795 vm.ctx.none()
796 }
797
798 #[pygetset]
799 fn errors(_zelf: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
800 vm.ctx.none()
801 }
802 }
803
804 #[derive(FromArgs, Clone)]
805 struct BufferSize {
806 #[pyarg(any, optional)]
807 buffer_size: OptionalArg<isize>,
808 }
809
810 bitflags::bitflags! {
811 #[derive(Copy, Clone, Debug, PartialEq, Default)]
812 struct BufferedFlags: u8 {
813 const DETACHED = 1 << 0;
814 const WRITABLE = 1 << 1;
815 const READABLE = 1 << 2;
816 }
817 }
818
819 #[derive(Debug, Default)]
820 struct BufferedData {
821 raw: Option<PyObjectRef>,
822 flags: BufferedFlags,
823 abs_pos: Offset,
824 buffer: Vec<u8>,
825 pos: Offset,
826 raw_pos: Offset,
827 read_end: Offset,
828 write_pos: Offset,
829 write_end: Offset,
830 }
831
832 impl BufferedData {
833 fn check_init(&self, vm: &VirtualMachine) -> PyResult<&PyObject> {
834 if let Some(raw) = &self.raw {
835 Ok(raw)
836 } else {
837 let msg = if self.flags.contains(BufferedFlags::DETACHED) {
838 "raw stream has been detached"
839 } else {
840 "I/O operation on uninitialized object"
841 };
842 Err(vm.new_value_error(msg))
843 }
844 }
845
846 #[inline]
847 const fn writable(&self) -> bool {
848 self.flags.contains(BufferedFlags::WRITABLE)
849 }
850
851 #[inline]
852 const fn readable(&self) -> bool {
853 self.flags.contains(BufferedFlags::READABLE)
854 }
855
856 #[inline]
857 const fn valid_read(&self) -> bool {
858 self.readable() && self.read_end != -1
859 }
860
861 #[inline]
862 const fn valid_write(&self) -> bool {
863 self.writable() && self.write_end != -1
864 }
865
866 #[inline]
867 const fn raw_offset(&self) -> Offset {
868 if (self.valid_read() || self.valid_write()) && self.raw_pos >= 0 {
869 self.raw_pos - self.pos
870 } else {
871 0
872 }
873 }
874
875 #[inline]
876 const fn readahead(&self) -> Offset {
877 if self.valid_read() {
878 self.read_end - self.pos
879 } else {
880 0
881 }
882 }
883
884 const fn reset_read(&mut self) {
885 self.read_end = -1;
886 }
887
888 const fn reset_write(&mut self) {
889 self.write_pos = 0;
890 self.write_end = -1;
891 }
892
893 fn flush(&mut self, vm: &VirtualMachine) -> PyResult<()> {
894 if !self.valid_write() || self.write_pos == self.write_end {
895 self.reset_write();
896 return Ok(());
897 }
898
899 let rewind = self.raw_offset() + (self.pos - self.write_pos);
900 if rewind != 0 {
901 self.raw_seek(-rewind, 1, vm)?;
902 self.raw_pos -= rewind;
903 }
904
905 while self.write_pos < self.write_end {
906 let n =
907 self.raw_write(None, self.write_pos as usize..self.write_end as usize, vm)?;
908 let n = match n {
909 Some(n) => n,
910 None => {
911 return Err(vm.invoke_exception(
913 vm.ctx.exceptions.blocking_io_error.to_owned(),
914 vec![
915 vm.new_pyobj(EAGAIN),
916 vm.new_pyobj("write could not complete without blocking"),
917 vm.new_pyobj(0),
918 ],
919 )?);
920 }
921 };
922 self.write_pos += n as Offset;
923 self.raw_pos = self.write_pos;
924 vm.check_signals()?;
925 }
926
927 self.reset_write();
928
929 Ok(())
930 }
931
932 fn flush_rewind(&mut self, vm: &VirtualMachine) -> PyResult<()> {
933 self.flush(vm)?;
934 if self.readable() {
935 let res = self.raw_seek(-self.raw_offset(), 1, vm);
936 self.reset_read();
937 res?;
938 }
939 Ok(())
940 }
941
942 fn raw_seek(&mut self, pos: Offset, whence: i32, vm: &VirtualMachine) -> PyResult<Offset> {
943 let ret = vm.call_method(self.check_init(vm)?, "seek", (pos, whence))?;
944 let offset = get_offset(ret, vm)?;
945 if offset < 0 {
946 return Err(
947 vm.new_os_error(format!("Raw stream returned invalid position {offset}"))
948 );
949 }
950 self.abs_pos = offset;
951 Ok(offset)
952 }
953
954 fn seek(&mut self, target: Offset, whence: i32, vm: &VirtualMachine) -> PyResult<Offset> {
955 if matches!(whence, 0 | 1) && self.readable() {
956 let current = self.raw_tell_cache(vm)?;
957 let available = self.readahead();
958 if available > 0 {
959 let offset = if whence == 0 {
960 target - (current - self.raw_offset())
961 } else {
962 target
963 };
964 if offset >= -self.pos && offset <= available {
965 self.pos += offset;
966 let result = current - available + offset;
969 return Ok(if result < 0 { 0 } else { result });
970 }
971 }
972 }
973 if self.writable() {
975 self.flush(vm)?;
976 }
977 let target = if whence == 1 {
978 target - self.raw_offset()
979 } else {
980 target
981 };
982 let res = self.raw_seek(target, whence, vm);
983 self.raw_pos = -1;
984 if res.is_ok() && self.readable() {
985 self.reset_read();
986 }
987 res
988 }
989
990 fn raw_tell(&mut self, vm: &VirtualMachine) -> PyResult<Offset> {
991 let raw = self.check_init(vm)?;
992 let ret = vm.call_method(raw, "tell", ())?;
993 let offset = get_offset(ret, vm)?;
994 if offset < 0 {
995 return Err(
996 vm.new_os_error(format!("Raw stream returned invalid position {offset}"))
997 );
998 }
999 self.abs_pos = offset;
1000 Ok(offset)
1001 }
1002
1003 fn raw_tell_cache(&mut self, vm: &VirtualMachine) -> PyResult<Offset> {
1004 if self.abs_pos == -1 {
1005 self.raw_tell(vm)
1006 } else {
1007 Ok(self.abs_pos)
1008 }
1009 }
1010
1011 fn raw_write(
1013 &mut self,
1014 buf: Option<PyBuffer>,
1015 buf_range: Range<usize>,
1016 vm: &VirtualMachine,
1017 ) -> PyResult<Option<usize>> {
1018 let len = buf_range.len();
1019
1020 let (mem_obj, write_buf) = if let Some(buf) = buf {
1023 let mem_obj =
1024 PyMemoryView::from_buffer_range(buf, buf_range, vm)?.into_ref(&vm.ctx);
1025 (mem_obj, None)
1026 } else {
1027 let v = core::mem::take(&mut self.buffer);
1028 let wb = VecBuffer::from(v).into_ref(&vm.ctx);
1029 let mem_obj =
1030 PyMemoryView::from_buffer_range(wb.clone().into_pybuffer(true), buf_range, vm)?
1031 .into_ref(&vm.ctx);
1032 (mem_obj, Some(wb))
1033 };
1034
1035 let res = loop {
1037 let res = vm.call_method(self.raw.as_ref().unwrap(), "write", (mem_obj.clone(),));
1038 match trap_eintr(res, vm) {
1039 Ok(Some(val)) => break Ok(val),
1040 Ok(None) => continue,
1041 Err(e) => break Err(e),
1042 }
1043 };
1044
1045 if let Some(wb) = write_buf {
1047 mem_obj.release();
1048 self.buffer = wb.take();
1049 }
1050
1051 let res = res?;
1052
1053 if vm.is_none(&res) {
1054 return Ok(None);
1055 }
1056 let n = isize::try_from_object(vm, res)?;
1057 if n < 0 || n as usize > len {
1058 return Err(vm.new_os_error(format!(
1059 "raw write() returned invalid length {n} (should have been between 0 and {len})"
1060 )));
1061 }
1062 if self.abs_pos != -1 {
1063 self.abs_pos += n as Offset
1064 }
1065 Ok(Some(n as usize))
1066 }
1067
1068 fn write(&mut self, obj: ArgBytesLike, vm: &VirtualMachine) -> PyResult<usize> {
1069 if !self.valid_read() && !self.valid_write() {
1070 self.pos = 0;
1071 self.raw_pos = 0;
1072 }
1073 let avail = self.buffer.len() - self.pos as usize;
1074 let buf_len;
1075 {
1076 let buf = obj.borrow_buf();
1077 buf_len = buf.len();
1078 if buf.len() <= avail {
1079 self.buffer[self.pos as usize..][..buf.len()].copy_from_slice(&buf);
1080 if !self.valid_write() || self.write_pos > self.pos {
1081 self.write_pos = self.pos
1082 }
1083 self.adjust_position(self.pos + buf.len() as Offset);
1084 if self.pos > self.write_end {
1085 self.write_end = self.pos
1086 }
1087 return Ok(buf.len());
1088 }
1089 }
1090
1091 match self.flush(vm) {
1094 Ok(()) => {}
1095 Err(e) if e.fast_isinstance(vm.ctx.exceptions.blocking_io_error) => {
1096 if self.readable() {
1097 self.reset_read();
1098 }
1099 let shift = self.write_pos;
1101 if shift > 0 {
1102 self.buffer
1103 .copy_within(shift as usize..self.write_end as usize, 0);
1104 self.write_end -= shift;
1105 self.raw_pos -= shift;
1106 self.pos -= shift;
1107 self.write_pos = 0;
1108 }
1109 let avail = self.buffer.len() - self.write_end as usize;
1110 if buf_len <= avail {
1111 let buf = obj.borrow_buf();
1113 self.buffer[self.write_end as usize..][..buf_len].copy_from_slice(&buf);
1114 self.write_end += buf_len as Offset;
1115 self.pos += buf_len as Offset;
1116 return Ok(buf_len);
1117 }
1118 let buf = obj.borrow_buf();
1120 self.buffer[self.write_end as usize..][..avail].copy_from_slice(&buf[..avail]);
1121 self.write_end += avail as Offset;
1122 self.pos += avail as Offset;
1123 return Err(vm.invoke_exception(
1124 vm.ctx.exceptions.blocking_io_error.to_owned(),
1125 vec![
1126 vm.new_pyobj(EAGAIN),
1127 vm.new_pyobj("write could not complete without blocking"),
1128 vm.new_pyobj(avail),
1129 ],
1130 )?);
1131 }
1132 Err(e) => return Err(e),
1133 }
1134
1135 let offset = self.raw_offset();
1137 if offset != 0 {
1138 self.raw_seek(-offset, 1, vm)?;
1139 self.raw_pos -= offset;
1140 }
1141
1142 let mut remaining = buf_len;
1143 let mut written = 0;
1144 let buffer: PyBuffer = obj.into();
1145 while remaining > self.buffer.len() {
1146 let res = self.raw_write(Some(buffer.clone()), written..buf_len, vm)?;
1147 match res {
1148 Some(n) => {
1149 written += n;
1150 if let Some(r) = remaining.checked_sub(n) {
1151 remaining = r
1152 } else {
1153 break;
1154 }
1155 vm.check_signals()?;
1156 }
1157 None => {
1158 if remaining > self.buffer.len() {
1160 let buf = buffer.as_contiguous().unwrap();
1162 let buffer_len = self.buffer.len();
1163 self.buffer.copy_from_slice(&buf[written..][..buffer_len]);
1164 self.raw_pos = 0;
1165 let buffer_size = self.buffer.len() as _;
1166 self.adjust_position(buffer_size);
1167 self.write_end = buffer_size;
1168 let chars_written = written + buffer_len;
1170 return Err(vm.invoke_exception(
1171 vm.ctx.exceptions.blocking_io_error.to_owned(),
1172 vec![
1173 vm.new_pyobj(EAGAIN),
1174 vm.new_pyobj("write could not complete without blocking"),
1175 vm.new_pyobj(chars_written),
1176 ],
1177 )?);
1178 } else {
1179 break;
1180 }
1181 }
1182 }
1183 }
1184 if self.readable() {
1185 self.reset_read();
1186 }
1187 if remaining > 0 {
1188 let buf = buffer.as_contiguous().unwrap();
1189 self.buffer[..remaining].copy_from_slice(&buf[written..][..remaining]);
1190 written += remaining;
1191 }
1192 self.write_pos = 0;
1193 self.write_end = remaining as _;
1194 self.adjust_position(remaining as _);
1195 self.raw_pos = 0;
1196
1197 Ok(written)
1198 }
1199
1200 fn active_read_slice(&self) -> &[u8] {
1201 &self.buffer[self.pos as usize..][..self.readahead() as usize]
1202 }
1203
1204 fn read_fast(&mut self, n: usize) -> Option<Vec<u8>> {
1205 let ret = self.active_read_slice().get(..n)?.to_vec();
1206 self.pos += n as Offset;
1207 Some(ret)
1208 }
1209
1210 fn read_generic(&mut self, n: usize, vm: &VirtualMachine) -> PyResult<Option<Vec<u8>>> {
1211 if let Some(fast) = self.read_fast(n) {
1212 return Ok(Some(fast));
1213 }
1214
1215 let current_size = self.readahead() as usize;
1216
1217 let mut out = vec![0u8; n];
1218 let mut remaining = n;
1219 let mut written = 0;
1220 if current_size > 0 {
1221 let slice = self.active_read_slice();
1222 out[..slice.len()].copy_from_slice(slice);
1223 remaining -= current_size;
1224 written += current_size;
1225 self.pos += current_size as Offset;
1226 }
1227 if self.writable() {
1228 self.flush_rewind(vm)?;
1229 }
1230 self.reset_read();
1231 macro_rules! handle_opt_read {
1232 ($x:expr) => {
1233 match ($x, written > 0) {
1234 (Some(0), _) | (None, true) => {
1235 out.truncate(written);
1236 return Ok(Some(out));
1237 }
1238 (Some(r), _) => r,
1239 (None, _) => return Ok(None),
1240 }
1241 };
1242 }
1243 while remaining > 0 && !self.buffer.is_empty() {
1244 let r = self.buffer.len() * (remaining / self.buffer.len());
1246 if r == 0 {
1247 break;
1248 }
1249 let r = self.raw_read(Either::A(Some(&mut out)), written..written + r, vm)?;
1250 let r = handle_opt_read!(r);
1251 remaining -= r;
1252 written += r;
1253 }
1254 self.pos = 0;
1255 self.raw_pos = 0;
1256 self.read_end = 0;
1257
1258 while remaining > 0 && (self.read_end as usize) < self.buffer.len() {
1259 let r = handle_opt_read!(self.fill_buffer(vm)?);
1260 if remaining > r {
1261 out[written..][..r].copy_from_slice(&self.buffer[self.pos as usize..][..r]);
1262 written += r;
1263 self.pos += r as Offset;
1264 remaining -= r;
1265 } else if remaining > 0 {
1266 out[written..][..remaining]
1267 .copy_from_slice(&self.buffer[self.pos as usize..][..remaining]);
1268 written += remaining;
1269 self.pos += remaining as Offset;
1270 remaining = 0;
1271 }
1272 if remaining == 0 {
1273 break;
1274 }
1275 }
1276
1277 Ok(Some(out))
1278 }
1279
1280 fn fill_buffer(&mut self, vm: &VirtualMachine) -> PyResult<Option<usize>> {
1281 let start = if self.valid_read() {
1282 self.read_end as usize
1283 } else {
1284 0
1285 };
1286 let buf_end = self.buffer.len();
1287 let res = self.raw_read(Either::A(None), start..buf_end, vm)?;
1288 if let Some(n) = res.filter(|n| *n > 0) {
1289 let new_start = (start + n) as Offset;
1290 self.read_end = new_start;
1291 self.raw_pos = new_start;
1292 }
1293 Ok(res)
1294 }
1295
1296 fn raw_read(
1297 &mut self,
1298 v: Either<Option<&mut Vec<u8>>, PyBuffer>,
1299 buf_range: Range<usize>,
1300 vm: &VirtualMachine,
1301 ) -> PyResult<Option<usize>> {
1302 let len = buf_range.len();
1303 let res = match v {
1304 Either::A(v) => {
1305 let v = v.unwrap_or(&mut self.buffer);
1306 let read_buf = VecBuffer::from(core::mem::take(v)).into_ref(&vm.ctx);
1307 let mem_obj = PyMemoryView::from_buffer_range(
1308 read_buf.clone().into_pybuffer(false),
1309 buf_range,
1310 vm,
1311 )?
1312 .into_ref(&vm.ctx);
1313
1314 let res = loop {
1316 let res = vm.call_method(
1317 self.raw.as_ref().unwrap(),
1318 "readinto",
1319 (mem_obj.clone(),),
1320 );
1321 match trap_eintr(res, vm) {
1322 Ok(Some(val)) => break Ok(val),
1323 Ok(None) => continue, Err(e) => break Err(e),
1325 }
1326 };
1327
1328 mem_obj.release();
1329 *v = read_buf.take();
1331
1332 res?
1333 }
1334 Either::B(buf) => {
1335 let mem_obj =
1336 PyMemoryView::from_buffer_range(buf, buf_range, vm)?.into_ref(&vm.ctx);
1337 loop {
1339 let res = vm.call_method(
1340 self.raw.as_ref().unwrap(),
1341 "readinto",
1342 (mem_obj.clone(),),
1343 );
1344 match trap_eintr(res, vm)? {
1345 Some(val) => break val,
1346 None => continue,
1347 }
1348 }
1349 }
1350 };
1351
1352 if vm.is_none(&res) {
1353 return Ok(None);
1354 }
1355 let (n, type_error) = match isize::try_from_object(vm, res.clone()) {
1357 Ok(n) => (n, None),
1358 Err(e) => (-1, Some(e)),
1359 };
1360 if n < 0 || n as usize > len {
1361 let os_error = vm.new_os_error(format!(
1362 "raw readinto() returned invalid length {n} (should have been between 0 and {len})"
1363 ));
1364 if let Some(cause) = type_error {
1365 os_error.set___cause__(Some(cause));
1366 }
1367 return Err(os_error);
1368 }
1369 if n > 0 && self.abs_pos != -1 {
1370 self.abs_pos += n as Offset
1371 }
1372 Ok(Some(n as usize))
1373 }
1374
1375 fn read_all(&mut self, vm: &VirtualMachine) -> PyResult<Option<PyBytesRef>> {
1376 let buf = self.active_read_slice();
1377 let data = if buf.is_empty() {
1378 None
1379 } else {
1380 let b = buf.to_vec();
1381 self.pos += buf.len() as Offset;
1382 Some(b)
1383 };
1384
1385 if self.writable() {
1386 self.flush_rewind(vm)?;
1387 }
1388
1389 let readall = vm
1390 .get_str_method(self.raw.clone().unwrap(), "readall")
1391 .transpose()?;
1392 if let Some(readall) = readall {
1393 let res = readall.call((), vm)?;
1394 let res = <Option<PyBytesRef>>::try_from_object(vm, res)?;
1395 let ret = if let Some(mut data) = data {
1396 if let Some(bytes) = res {
1397 data.extend_from_slice(bytes.as_bytes());
1398 }
1399 Some(PyBytes::from(data).into_ref(&vm.ctx))
1400 } else {
1401 res
1402 };
1403 return Ok(ret);
1404 }
1405
1406 let mut chunks = Vec::new();
1407
1408 let mut read_size = 0;
1409 loop {
1410 let read_data = loop {
1412 let res = vm.call_method(self.raw.as_ref().unwrap(), "read", ());
1413 match trap_eintr(res, vm)? {
1414 Some(val) => break val,
1415 None => continue,
1416 }
1417 };
1418 let read_data = <Option<PyBytesRef>>::try_from_object(vm, read_data)?;
1419
1420 match read_data {
1421 Some(b) if !b.as_bytes().is_empty() => {
1422 let l = b.as_bytes().len();
1423 read_size += l;
1424 if self.abs_pos != -1 {
1425 self.abs_pos += l as Offset;
1426 }
1427 chunks.push(b);
1428 }
1429 read_data => {
1430 let ret = if data.is_none() && read_size == 0 {
1431 read_data
1432 } else {
1433 let mut data = data.unwrap_or_default();
1434 data.reserve(read_size);
1435 for bytes in &chunks {
1436 data.extend_from_slice(bytes.as_bytes())
1437 }
1438 Some(PyBytes::from(data).into_ref(&vm.ctx))
1439 };
1440 break Ok(ret);
1441 }
1442 }
1443 }
1444 }
1445
1446 const fn adjust_position(&mut self, new_pos: Offset) {
1447 self.pos = new_pos;
1448 if self.valid_read() && self.read_end < self.pos {
1449 self.read_end = self.pos
1450 }
1451 }
1452
1453 fn peek(&mut self, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
1454 let have = self.readahead();
1455 let slice = if have > 0 {
1456 &self.buffer[self.pos as usize..][..have as usize]
1457 } else {
1458 self.reset_read();
1459 let r = self.fill_buffer(vm)?.unwrap_or(0);
1460 self.pos = 0;
1461 &self.buffer[..r]
1462 };
1463 Ok(slice.to_vec())
1464 }
1465
1466 fn readinto_generic(
1467 &mut self,
1468 buf: PyBuffer,
1469 readinto1: bool,
1470 vm: &VirtualMachine,
1471 ) -> PyResult<Option<usize>> {
1472 let mut written = 0;
1473 let n = self.readahead();
1474 let buf_len;
1475 {
1476 let mut b = buf.as_contiguous_mut().unwrap();
1477 buf_len = b.len();
1478 if n > 0 {
1479 if n as usize >= b.len() {
1480 b.copy_from_slice(&self.buffer[self.pos as usize..][..buf_len]);
1481 self.pos += buf_len as Offset;
1482 return Ok(Some(buf_len));
1483 }
1484 b[..n as usize]
1485 .copy_from_slice(&self.buffer[self.pos as usize..][..n as usize]);
1486 self.pos += n;
1487 written = n as usize;
1488 }
1489 }
1490 if self.writable() {
1491 self.flush_rewind(vm)?;
1492 }
1493 self.reset_read();
1494 self.pos = 0;
1495
1496 let mut remaining = buf_len - written;
1497 while remaining > 0 {
1498 let n = if remaining > self.buffer.len() {
1499 self.raw_read(Either::B(buf.clone()), written..written + remaining, vm)?
1500 } else if !(readinto1 && written != 0) {
1501 let n = self.fill_buffer(vm)?;
1502 if let Some(n) = n.filter(|&n| n > 0) {
1503 let n = core::cmp::min(n, remaining);
1504 buf.as_contiguous_mut().unwrap()[written..][..n]
1505 .copy_from_slice(&self.buffer[self.pos as usize..][..n]);
1506 self.pos += n as Offset;
1507 written += n;
1508 remaining -= n;
1509 continue;
1510 }
1511 n
1512 } else {
1513 break;
1514 };
1515 let n = match n {
1516 Some(0) => break,
1517 None if written > 0 => break,
1518 None => return Ok(None),
1519 Some(n) => n,
1520 };
1521
1522 if readinto1 {
1523 written += n;
1524 break;
1525 }
1526 written += n;
1527 remaining -= n;
1528 }
1529
1530 Ok(Some(written))
1531 }
1532 }
1533
1534 pub fn get_offset(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<Offset> {
1535 let int = obj.try_index(vm)?;
1536 int.as_bigint().try_into().map_err(|_| {
1537 vm.new_value_error(format!(
1538 "cannot fit '{}' into an offset-sized integer",
1539 obj.class().name()
1540 ))
1541 })
1542 }
1543
1544 pub fn repr_file_obj_name(obj: &PyObject, vm: &VirtualMachine) -> PyResult<Option<PyStrRef>> {
1545 let name = match obj.get_attr("name", vm) {
1546 Ok(name) => Some(name),
1547 Err(e)
1548 if e.fast_isinstance(vm.ctx.exceptions.attribute_error)
1549 || e.fast_isinstance(vm.ctx.exceptions.value_error) =>
1550 {
1551 None
1552 }
1553 Err(e) => return Err(e),
1554 };
1555 match name {
1556 Some(name) => {
1557 if let Some(_guard) = ReprGuard::enter(vm, obj) {
1558 name.repr(vm).map(Some)
1559 } else {
1560 Err(vm.new_runtime_error(format!(
1561 "reentrant call inside {}.__repr__",
1562 obj.class().slot_name()
1563 )))
1564 }
1565 }
1566 None => Ok(None),
1567 }
1568 }
1569
1570 #[pyclass]
1571 trait BufferedMixin: PyPayload + StaticType {
1572 const CLASS_NAME: &'static str;
1573 const READABLE: bool;
1574 const WRITABLE: bool;
1575 const SEEKABLE: bool = false;
1576
1577 fn data(&self) -> &PyThreadMutex<BufferedData>;
1578 fn closing(&self) -> &AtomicBool;
1579 fn finalizing(&self) -> &AtomicBool;
1580
1581 fn lock(&self, vm: &VirtualMachine) -> PyResult<PyThreadMutexGuard<'_, BufferedData>> {
1582 self.data()
1583 .lock_wrapped(|do_lock| vm.allow_threads(do_lock))
1584 .ok_or_else(|| vm.new_runtime_error("reentrant call inside buffered io"))
1585 }
1586
1587 #[pyslot]
1588 fn slot_init(zelf: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
1589 let zelf: PyRef<Self> = zelf.try_into_value(vm)?;
1590 let (raw, BufferSize { buffer_size }): (PyObjectRef, _) =
1591 args.bind(vm).map_err(|e| {
1592 let str_repr = e
1593 .__str__(vm)
1594 .as_ref()
1595 .map_or("<error getting exception str>".as_ref(), |s| s.as_wtf8())
1596 .to_owned();
1597 let msg = format!("{}() {}", Self::CLASS_NAME, str_repr);
1598 vm.new_exception_msg(e.class().to_owned(), msg.into())
1599 })?;
1600 zelf.init(raw, BufferSize { buffer_size }, vm)
1601 }
1602
1603 fn init(
1604 &self,
1605 raw: PyObjectRef,
1606 BufferSize { buffer_size }: BufferSize,
1607 vm: &VirtualMachine,
1608 ) -> PyResult<()> {
1609 let mut data = self.lock(vm)?;
1610 data.raw = None;
1611 data.flags.remove(BufferedFlags::DETACHED);
1612
1613 let buffer_size = match buffer_size {
1614 OptionalArg::Present(i) if i <= 0 => {
1615 return Err(vm.new_value_error("buffer size must be strictly positive"));
1616 }
1617 OptionalArg::Present(i) => i as usize,
1618 OptionalArg::Missing => DEFAULT_BUFFER_SIZE,
1619 };
1620
1621 if Self::SEEKABLE {
1622 check_seekable(&raw, vm)?;
1623 }
1624 if Self::READABLE {
1625 data.flags.insert(BufferedFlags::READABLE);
1626 check_readable(&raw, vm)?;
1627 }
1628 if Self::WRITABLE {
1629 data.flags.insert(BufferedFlags::WRITABLE);
1630 check_writable(&raw, vm)?;
1631 }
1632
1633 data.buffer = vec![0; buffer_size];
1634
1635 if Self::READABLE {
1636 data.reset_read();
1637 }
1638 if Self::WRITABLE {
1639 data.reset_write();
1640 }
1641 if Self::SEEKABLE {
1642 data.pos = 0;
1643 }
1644
1645 data.raw = Some(raw);
1646
1647 Ok(())
1648 }
1649
1650 #[pymethod]
1651 fn seek(
1652 &self,
1653 target: PyObjectRef,
1654 whence: OptionalArg<i32>,
1655 vm: &VirtualMachine,
1656 ) -> PyResult<Offset> {
1657 let whence = whence.unwrap_or(0);
1658 if !validate_whence(whence) {
1659 return Err(vm.new_value_error(format!("whence value {whence} unsupported")));
1660 }
1661 let mut data = self.lock(vm)?;
1662 let raw = data.check_init(vm)?;
1663 ensure_unclosed(raw, "seek of closed file", vm)?;
1664 check_seekable(raw, vm)?;
1665 let target = get_offset(target, vm)?;
1666 data.seek(target, whence, vm)
1667 }
1668
1669 #[pymethod]
1670 fn tell(&self, vm: &VirtualMachine) -> PyResult<Offset> {
1671 let mut data = self.lock(vm)?;
1672 let raw_tell = data.raw_tell(vm)?;
1673 let raw_offset = data.raw_offset();
1674 let mut pos = raw_tell - raw_offset;
1675 if pos < 0 {
1677 pos = 0;
1678 }
1679 Ok(pos)
1680 }
1681
1682 #[pymethod]
1683 fn truncate(
1684 zelf: PyRef<Self>,
1685 pos: OptionalOption<PyObjectRef>,
1686 vm: &VirtualMachine,
1687 ) -> PyResult {
1688 let pos = pos.flatten().to_pyobject(vm);
1689 let mut data = zelf.lock(vm)?;
1690 data.check_init(vm)?;
1691 if !data.writable() {
1692 return Err(new_unsupported_operation(vm, "truncate".to_owned()));
1693 }
1694 data.flush_rewind(vm)?;
1695 let res = vm.call_method(data.raw.as_ref().unwrap(), "truncate", (pos,))?;
1696 let _ = data.raw_tell(vm);
1697 Ok(res)
1698 }
1699 #[pymethod]
1700 fn detach(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
1701 vm.call_method(zelf.as_object(), "flush", ())?;
1702 let mut data = zelf.lock(vm)?;
1703 data.flags.insert(BufferedFlags::DETACHED);
1704 data.raw
1705 .take()
1706 .ok_or_else(|| vm.new_value_error("raw stream has been detached"))
1707 }
1708
1709 #[pymethod]
1710 fn seekable(&self, vm: &VirtualMachine) -> PyResult {
1711 vm.call_method(self.lock(vm)?.check_init(vm)?, "seekable", ())
1712 }
1713
1714 #[pygetset]
1715 fn raw(&self, vm: &VirtualMachine) -> PyResult<Option<PyObjectRef>> {
1716 Ok(self.lock(vm)?.raw.clone())
1717 }
1718
1719 fn get_raw_unlocked(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
1721 let data = self.lock(vm)?;
1722 Ok(data.check_init(vm)?.to_owned())
1723 }
1724
1725 #[pygetset]
1726 fn closed(&self, vm: &VirtualMachine) -> PyResult {
1727 self.get_raw_unlocked(vm)?.get_attr("closed", vm)
1728 }
1729
1730 #[pygetset]
1731 fn name(&self, vm: &VirtualMachine) -> PyResult {
1732 self.get_raw_unlocked(vm)?.get_attr("name", vm)
1733 }
1734
1735 #[pygetset]
1736 fn mode(&self, vm: &VirtualMachine) -> PyResult {
1737 self.get_raw_unlocked(vm)?.get_attr("mode", vm)
1738 }
1739
1740 #[pymethod]
1741 fn fileno(&self, vm: &VirtualMachine) -> PyResult {
1742 vm.call_method(self.lock(vm)?.check_init(vm)?, "fileno", ())
1743 }
1744
1745 #[pymethod]
1746 fn isatty(&self, vm: &VirtualMachine) -> PyResult {
1747 vm.call_method(self.lock(vm)?.check_init(vm)?, "isatty", ())
1748 }
1749
1750 #[pyslot]
1751 fn slot_repr(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<PyRef<PyStr>> {
1752 let name_repr = repr_file_obj_name(zelf, vm)?;
1753 let cls = zelf.class();
1754 let slot_name = cls.slot_name();
1755 let repr = if let Some(name_repr) = name_repr {
1756 format!("<{slot_name} name={name_repr}>")
1757 } else {
1758 format!("<{slot_name}>")
1759 };
1760 Ok(vm.ctx.new_str(repr))
1761 }
1762
1763 #[pymethod]
1764 fn __repr__(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyRef<PyStr>> {
1765 Self::slot_repr(&zelf, vm)
1766 }
1767
1768 fn close_strict(&self, vm: &VirtualMachine) -> PyResult {
1769 let mut data = self.lock(vm)?;
1770 let raw = data.check_init(vm)?;
1771 if file_closed(raw, vm)? {
1772 return Ok(vm.ctx.none());
1773 }
1774 let flush_res = data.flush(vm);
1775 let close_res = vm.call_method(data.raw.as_ref().unwrap(), "close", ());
1776 exception_chain(flush_res, close_res)
1777 }
1778
1779 #[pymethod]
1780 fn close(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
1781 let raw = {
1783 let data = zelf.lock(vm)?;
1784 let raw = data.check_init(vm)?;
1785 if file_closed(raw, vm)? {
1786 return Ok(vm.ctx.none());
1787 }
1788 raw.to_owned()
1789 };
1790 if zelf.finalizing().load(Ordering::Relaxed) {
1791 let _ = vm.call_method(&raw, "_dealloc_warn", (zelf.as_object().to_owned(),));
1793 }
1794 zelf.closing().store(true, Ordering::Release);
1796 let flush_res = vm.call_method(zelf.as_object(), "flush", ()).map(drop);
1797 let close_res = vm.call_method(&raw, "close", ());
1798 exception_chain(flush_res, close_res)
1799 }
1800
1801 #[pymethod]
1802 fn readable(&self) -> bool {
1803 Self::READABLE
1804 }
1805
1806 #[pymethod]
1807 fn writable(&self) -> bool {
1808 Self::WRITABLE
1809 }
1810
1811 #[pymethod]
1812 fn __getstate__(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
1813 Err(vm.new_type_error(format!("cannot pickle '{}' instances", zelf.class().name())))
1814 }
1815
1816 #[pymethod]
1817 fn __reduce_ex__(zelf: PyObjectRef, proto: usize, vm: &VirtualMachine) -> PyResult {
1818 if zelf.class().is(Self::static_type()) {
1819 return Err(
1820 vm.new_type_error(format!("cannot pickle '{}' object", zelf.class().name()))
1821 );
1822 }
1823 let _ = proto;
1824 reduce_ex_for_subclass(zelf, vm)
1825 }
1826
1827 #[pymethod]
1828 fn _dealloc_warn(
1829 zelf: PyRef<Self>,
1830 source: PyObjectRef,
1831 vm: &VirtualMachine,
1832 ) -> PyResult<()> {
1833 let raw = {
1835 let data = zelf.lock(vm)?;
1836 data.raw.clone()
1837 };
1838 if let Some(raw) = raw {
1839 let _ = vm.call_method(&raw, "_dealloc_warn", (source,));
1840 }
1841 Ok(())
1842 }
1843 }
1844
1845 #[pyclass]
1846 trait BufferedReadable: PyPayload {
1847 type Reader: BufferedMixin;
1848
1849 fn reader(&self) -> &Self::Reader;
1850
1851 #[pymethod]
1852 fn read(&self, size: OptionalSize, vm: &VirtualMachine) -> PyResult<Option<PyBytesRef>> {
1853 let mut data = self.reader().lock(vm)?;
1854 let raw = data.check_init(vm)?;
1855 let n = size.size.map(|s| *s).unwrap_or(-1);
1856 if n < -1 {
1857 return Err(vm.new_value_error("read length must be non-negative or -1"));
1858 }
1859 ensure_unclosed(raw, "read of closed file", vm)?;
1860 match n.to_usize() {
1861 Some(n) => data
1862 .read_generic(n, vm)
1863 .map(|x| x.map(|b| PyBytes::from(b).into_ref(&vm.ctx))),
1864 None => data.read_all(vm),
1865 }
1866 }
1867
1868 #[pymethod]
1869 fn peek(&self, _size: OptionalSize, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
1870 let mut data = self.reader().lock(vm)?;
1871 let raw = data.check_init(vm)?;
1872 ensure_unclosed(raw, "peek of closed file", vm)?;
1873
1874 if data.writable() {
1875 let _ = data.flush_rewind(vm);
1876 }
1877 data.peek(vm)
1878 }
1879
1880 #[pymethod]
1881 fn read1(&self, size: OptionalSize, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
1882 let mut data = self.reader().lock(vm)?;
1883 let raw = data.check_init(vm)?;
1884 ensure_unclosed(raw, "read of closed file", vm)?;
1885 let n = size.to_usize().unwrap_or(data.buffer.len());
1886 if n == 0 {
1887 return Ok(Vec::new());
1888 }
1889 let have = data.readahead();
1890 if have > 0 {
1891 let n = core::cmp::min(have as usize, n);
1892 return Ok(data.read_fast(n).unwrap());
1893 }
1894 if data.writable() {
1896 data.flush_rewind(vm)?;
1897 }
1898 let mut v = vec![0; n];
1899 data.reset_read();
1900 let r = data
1901 .raw_read(Either::A(Some(&mut v)), 0..n, vm)?
1902 .unwrap_or(0);
1903 v.truncate(r);
1904 v.shrink_to_fit();
1905 Ok(v)
1906 }
1907
1908 #[pymethod]
1909 fn readinto(&self, buf: ArgMemoryBuffer, vm: &VirtualMachine) -> PyResult<Option<usize>> {
1910 let mut data = self.reader().lock(vm)?;
1911 let raw = data.check_init(vm)?;
1912 ensure_unclosed(raw, "readinto of closed file", vm)?;
1913 data.readinto_generic(buf.into(), false, vm)
1914 }
1915
1916 #[pymethod]
1917 fn readinto1(&self, buf: ArgMemoryBuffer, vm: &VirtualMachine) -> PyResult<Option<usize>> {
1918 let mut data = self.reader().lock(vm)?;
1919 let raw = data.check_init(vm)?;
1920 ensure_unclosed(raw, "readinto of closed file", vm)?;
1921 data.readinto_generic(buf.into(), true, vm)
1922 }
1923
1924 #[pymethod]
1925 fn flush(&self, vm: &VirtualMachine) -> PyResult<()> {
1926 let raw = {
1929 let data = self.reader().lock(vm)?;
1930 data.check_init(vm)?.to_owned()
1931 };
1932 ensure_unclosed(&raw, "flush of closed file", vm)?;
1933 vm.call_method(&raw, "flush", ())?;
1934 Ok(())
1935 }
1936 }
1937
1938 fn exception_chain<T>(e1: PyResult<()>, e2: PyResult<T>) -> PyResult<T> {
1939 match (e1, e2) {
1940 (Err(e1), Err(e)) => {
1941 e.set___context__(Some(e1));
1942 Err(e)
1943 }
1944 (Err(e), Ok(_)) | (Ok(()), Err(e)) => Err(e),
1945 (Ok(()), Ok(close_res)) => Ok(close_res),
1946 }
1947 }
1948
1949 #[pyattr]
1950 #[pyclass(name = "BufferedReader", base = _BufferedIOBase)]
1951 #[derive(Debug, Default)]
1952 struct BufferedReader {
1953 _base: _BufferedIOBase,
1954 data: PyThreadMutex<BufferedData>,
1955 closing: AtomicBool,
1956 finalizing: AtomicBool,
1957 }
1958
1959 impl BufferedMixin for BufferedReader {
1960 const CLASS_NAME: &'static str = "BufferedReader";
1961 const READABLE: bool = true;
1962 const WRITABLE: bool = false;
1963
1964 fn data(&self) -> &PyThreadMutex<BufferedData> {
1965 &self.data
1966 }
1967
1968 fn closing(&self) -> &AtomicBool {
1969 &self.closing
1970 }
1971
1972 fn finalizing(&self) -> &AtomicBool {
1973 &self.finalizing
1974 }
1975 }
1976
1977 impl BufferedReadable for BufferedReader {
1978 type Reader = Self;
1979
1980 fn reader(&self) -> &Self::Reader {
1981 self
1982 }
1983 }
1984
1985 #[pyclass(
1986 with(Constructor, BufferedMixin, BufferedReadable, Destructor),
1987 flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
1988 )]
1989 impl BufferedReader {}
1990
1991 impl Destructor for BufferedReader {
1992 fn slot_del(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
1993 if let Some(buf) = zelf.downcast_ref::<BufferedReader>() {
1994 buf.finalizing.store(true, Ordering::Relaxed);
1995 }
1996 iobase_finalize(zelf, vm);
1997 Ok(())
1998 }
1999
2000 #[cold]
2001 fn del(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<()> {
2002 unreachable!("slot_del is implemented")
2003 }
2004 }
2005
2006 impl DefaultConstructor for BufferedReader {}
2007
2008 #[pyclass]
2009 trait BufferedWritable: PyPayload {
2010 type Writer: BufferedMixin;
2011
2012 fn writer(&self) -> &Self::Writer;
2013
2014 #[pymethod]
2015 fn write(&self, obj: ArgBytesLike, vm: &VirtualMachine) -> PyResult<usize> {
2016 if self.writer().closing().load(Ordering::Acquire) {
2021 loop {
2022 let raw = {
2023 let data = self.writer().lock(vm)?;
2024 match &data.raw {
2025 Some(raw) => raw.to_owned(),
2026 None => break, }
2028 };
2029 if file_closed(&raw, vm)? {
2030 break;
2031 }
2032 std::thread::yield_now();
2034 }
2035 return Err(vm.new_value_error("write to closed file"));
2036 }
2037 let mut data = self.writer().lock(vm)?;
2038 let raw = data.check_init(vm)?;
2039 ensure_unclosed(raw, "write to closed file", vm)?;
2040
2041 data.write(obj, vm)
2042 }
2043
2044 #[pymethod]
2045 fn flush(&self, vm: &VirtualMachine) -> PyResult<()> {
2046 let mut data = self.writer().lock(vm)?;
2047 let raw = data.check_init(vm)?;
2048 ensure_unclosed(raw, "flush of closed file", vm)?;
2049 data.flush_rewind(vm)
2050 }
2051 }
2052
2053 #[pyattr]
2054 #[pyclass(name = "BufferedWriter", base = _BufferedIOBase)]
2055 #[derive(Debug, Default)]
2056 struct BufferedWriter {
2057 _base: _BufferedIOBase,
2058 data: PyThreadMutex<BufferedData>,
2059 closing: AtomicBool,
2060 finalizing: AtomicBool,
2061 }
2062
2063 impl BufferedMixin for BufferedWriter {
2064 const CLASS_NAME: &'static str = "BufferedWriter";
2065 const READABLE: bool = false;
2066 const WRITABLE: bool = true;
2067
2068 fn data(&self) -> &PyThreadMutex<BufferedData> {
2069 &self.data
2070 }
2071
2072 fn closing(&self) -> &AtomicBool {
2073 &self.closing
2074 }
2075
2076 fn finalizing(&self) -> &AtomicBool {
2077 &self.finalizing
2078 }
2079 }
2080
2081 impl BufferedWritable for BufferedWriter {
2082 type Writer = Self;
2083
2084 fn writer(&self) -> &Self::Writer {
2085 self
2086 }
2087 }
2088
2089 #[pyclass(
2090 with(Constructor, BufferedMixin, BufferedWritable, Destructor),
2091 flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
2092 )]
2093 impl BufferedWriter {}
2094
2095 impl Destructor for BufferedWriter {
2096 fn slot_del(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
2097 if let Some(buf) = zelf.downcast_ref::<BufferedWriter>() {
2098 buf.finalizing.store(true, Ordering::Relaxed);
2099 }
2100 iobase_finalize(zelf, vm);
2101 Ok(())
2102 }
2103
2104 #[cold]
2105 fn del(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<()> {
2106 unreachable!("slot_del is implemented")
2107 }
2108 }
2109
2110 impl DefaultConstructor for BufferedWriter {}
2111
2112 #[pyattr]
2113 #[pyclass(name = "BufferedRandom", base = _BufferedIOBase)]
2114 #[derive(Debug, Default)]
2115 struct BufferedRandom {
2116 _base: _BufferedIOBase,
2117 data: PyThreadMutex<BufferedData>,
2118 closing: AtomicBool,
2119 finalizing: AtomicBool,
2120 }
2121
2122 impl BufferedMixin for BufferedRandom {
2123 const CLASS_NAME: &'static str = "BufferedRandom";
2124 const READABLE: bool = true;
2125 const WRITABLE: bool = true;
2126 const SEEKABLE: bool = true;
2127
2128 fn data(&self) -> &PyThreadMutex<BufferedData> {
2129 &self.data
2130 }
2131
2132 fn closing(&self) -> &AtomicBool {
2133 &self.closing
2134 }
2135
2136 fn finalizing(&self) -> &AtomicBool {
2137 &self.finalizing
2138 }
2139 }
2140
2141 impl BufferedReadable for BufferedRandom {
2142 type Reader = Self;
2143
2144 fn reader(&self) -> &Self::Reader {
2145 self
2146 }
2147 }
2148
2149 impl BufferedWritable for BufferedRandom {
2150 type Writer = Self;
2151
2152 fn writer(&self) -> &Self::Writer {
2153 self
2154 }
2155 }
2156
2157 #[pyclass(
2158 with(
2159 Constructor,
2160 BufferedMixin,
2161 BufferedReadable,
2162 BufferedWritable,
2163 Destructor
2164 ),
2165 flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
2166 )]
2167 impl BufferedRandom {}
2168
2169 impl Destructor for BufferedRandom {
2170 fn slot_del(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
2171 if let Some(buf) = zelf.downcast_ref::<BufferedRandom>() {
2172 buf.finalizing.store(true, Ordering::Relaxed);
2173 }
2174 iobase_finalize(zelf, vm);
2175 Ok(())
2176 }
2177
2178 #[cold]
2179 fn del(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<()> {
2180 unreachable!("slot_del is implemented")
2181 }
2182 }
2183
2184 impl DefaultConstructor for BufferedRandom {}
2185
2186 #[pyattr]
2187 #[pyclass(name = "BufferedRWPair", base = _BufferedIOBase)]
2188 #[derive(Debug, Default)]
2189 struct BufferedRWPair {
2190 _base: _BufferedIOBase,
2191 read: BufferedReader,
2192 write: BufferedWriter,
2193 }
2194
2195 impl BufferedReadable for BufferedRWPair {
2196 type Reader = BufferedReader;
2197
2198 fn reader(&self) -> &Self::Reader {
2199 &self.read
2200 }
2201 }
2202
2203 impl BufferedWritable for BufferedRWPair {
2204 type Writer = BufferedWriter;
2205
2206 fn writer(&self) -> &Self::Writer {
2207 &self.write
2208 }
2209 }
2210
2211 impl DefaultConstructor for BufferedRWPair {}
2212
2213 impl Initializer for BufferedRWPair {
2214 type Args = (PyObjectRef, PyObjectRef, BufferSize);
2215
2216 fn init(
2217 zelf: PyRef<Self>,
2218 (reader, writer, buffer_size): Self::Args,
2219 vm: &VirtualMachine,
2220 ) -> PyResult<()> {
2221 zelf.read.init(reader, buffer_size.clone(), vm)?;
2222 zelf.write.init(writer, buffer_size, vm)?;
2223 Ok(())
2224 }
2225 }
2226
2227 #[pyclass(
2228 with(
2229 Constructor,
2230 Initializer,
2231 BufferedReadable,
2232 BufferedWritable,
2233 Destructor
2234 ),
2235 flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
2236 )]
2237 impl BufferedRWPair {
2238 #[pymethod]
2239 fn flush(&self, vm: &VirtualMachine) -> PyResult<()> {
2240 self.write.flush(vm)
2241 }
2242
2243 #[pymethod]
2244 const fn readable(&self) -> bool {
2245 true
2246 }
2247 #[pymethod]
2248 const fn writable(&self) -> bool {
2249 true
2250 }
2251
2252 #[pygetset]
2253 fn closed(&self, vm: &VirtualMachine) -> PyResult {
2254 self.write.closed(vm)
2255 }
2256
2257 #[pymethod]
2258 fn isatty(&self, vm: &VirtualMachine) -> PyResult {
2259 let res = self.read.isatty(vm)?;
2261 if res.clone().try_to_bool(vm)? {
2262 Ok(res)
2263 } else {
2264 self.write.isatty(vm)
2265 }
2266 }
2267
2268 #[pymethod]
2269 fn close(&self, vm: &VirtualMachine) -> PyResult {
2270 let write_res = self.write.close_strict(vm).map(drop);
2271 let read_res = self.read.close_strict(vm);
2272 exception_chain(write_res, read_res)
2273 }
2274 }
2275
2276 impl Destructor for BufferedRWPair {
2277 fn slot_del(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
2278 iobase_finalize(zelf, vm);
2279 Ok(())
2280 }
2281
2282 #[cold]
2283 fn del(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<()> {
2284 unreachable!("slot_del is implemented")
2285 }
2286 }
2287
2288 #[derive(FromArgs)]
2289 struct TextIOWrapperArgs {
2290 #[pyarg(any, default)]
2291 encoding: Option<PyUtf8StrRef>,
2292 #[pyarg(any, default)]
2293 errors: Option<PyUtf8StrRef>,
2294 #[pyarg(any, default)]
2295 newline: OptionalOption<Newlines>,
2296 #[pyarg(any, default)]
2297 line_buffering: OptionalOption<PyObjectRef>,
2298 #[pyarg(any, default)]
2299 write_through: OptionalOption<PyObjectRef>,
2300 }
2301
2302 #[derive(Debug, Copy, Clone, Default, PartialEq)]
2303 enum Newlines {
2304 #[default]
2305 Universal,
2306 Passthrough,
2307 Lf,
2308 Cr,
2309 Crlf,
2310 }
2311
2312 impl Newlines {
2313 fn find_newline(&self, s: &Wtf8) -> Result<usize, usize> {
2316 let len = s.len();
2317 match self {
2318 Self::Universal | Self::Lf => s.find("\n".as_ref()).map(|p| p + 1).ok_or(len),
2319 Self::Passthrough => {
2320 let bytes = s.as_bytes();
2321 memchr::memchr2(b'\n', b'\r', bytes)
2322 .map(|p| {
2323 let nl_len =
2324 if bytes[p] == b'\r' && bytes.get(p + 1).copied() == Some(b'\n') {
2325 2
2326 } else {
2327 1
2328 };
2329 p + nl_len
2330 })
2331 .ok_or(len)
2332 }
2333 Self::Cr => s.find("\r".as_ref()).map(|p| p + 1).ok_or(len),
2334 Self::Crlf => {
2335 let mut searched = 0;
2337 let mut remaining = s.as_bytes();
2338 loop {
2339 match memchr::memchr(b'\r', remaining) {
2340 Some(p) => match remaining.get(p + 1) {
2341 Some(&ch_after_cr) => {
2342 let pos_after = p + 2;
2343 if ch_after_cr == b'\n' {
2344 break Ok(searched + pos_after);
2345 } else {
2346 searched += pos_after;
2347 remaining = &remaining[pos_after..];
2348 continue;
2349 }
2350 }
2351 None => break Err(searched + p),
2352 },
2353 None => break Err(len),
2354 }
2355 }
2356 }
2357 }
2358 }
2359 }
2360
2361 impl TryFromObject for Newlines {
2362 fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
2363 let nl = if vm.is_none(&obj) {
2364 Self::Universal
2365 } else {
2366 let s = obj.downcast::<PyStr>().map_err(|obj| {
2367 vm.new_type_error(format!(
2368 "newline argument must be str or None, not {}",
2369 obj.class().name()
2370 ))
2371 })?;
2372 let wtf8 = s.as_wtf8();
2373 if !wtf8.is_utf8() {
2374 let repr = s.repr(vm)?.as_str().to_owned();
2375 return Err(vm.new_value_error(format!("illegal newline value: {repr}")));
2376 }
2377 let s_str = wtf8.as_str().expect("checked utf8");
2378 match s_str {
2379 "" => Self::Passthrough,
2380 "\n" => Self::Lf,
2381 "\r" => Self::Cr,
2382 "\r\n" => Self::Crlf,
2383 _ => return Err(vm.new_value_error(format!("illegal newline value: {s}"))),
2384 }
2385 };
2386 Ok(nl)
2387 }
2388 }
2389
2390 fn reduce_ex_for_subclass(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
2391 let cls = zelf.class();
2392 let new = vm
2393 .get_attribute_opt(cls.to_owned().into(), "__new__")?
2394 .ok_or_else(|| vm.new_attribute_error("type has no attribute '__new__'"))?;
2395 let args = vm.ctx.new_tuple(vec![cls.to_owned().into()]);
2396 let state = if let Some(getstate) = vm.get_attribute_opt(zelf.clone(), "__getstate__")? {
2397 getstate.call((), vm)?
2398 } else if let Ok(dict) = zelf.get_attr("__dict__", vm) {
2399 dict
2400 } else {
2401 vm.ctx.none()
2402 };
2403 Ok(vm.ctx.new_tuple(vec![new, args.into(), state]).into())
2404 }
2405
2406 #[derive(Debug, Default, Copy, Clone)]
2408 struct Utf8size {
2409 bytes: usize,
2410 chars: usize,
2411 }
2412
2413 impl Utf8size {
2414 fn len_pystr(s: &PyStr) -> Self {
2415 Self {
2416 bytes: s.byte_len(),
2417 chars: s.char_len(),
2418 }
2419 }
2420
2421 fn len_str(s: &Wtf8) -> Self {
2422 Self {
2423 bytes: s.len(),
2424 chars: s.code_points().count(),
2425 }
2426 }
2427 }
2428
2429 impl core::ops::Add for Utf8size {
2430 type Output = Self;
2431
2432 #[inline]
2433 fn add(mut self, rhs: Self) -> Self {
2434 self += rhs;
2435 self
2436 }
2437 }
2438
2439 impl core::ops::AddAssign for Utf8size {
2440 #[inline]
2441 fn add_assign(&mut self, rhs: Self) {
2442 self.bytes += rhs.bytes;
2443 self.chars += rhs.chars;
2444 }
2445 }
2446
2447 impl core::ops::Sub for Utf8size {
2448 type Output = Self;
2449
2450 #[inline]
2451 fn sub(mut self, rhs: Self) -> Self {
2452 self -= rhs;
2453 self
2454 }
2455 }
2456
2457 impl core::ops::SubAssign for Utf8size {
2458 #[inline]
2459 fn sub_assign(&mut self, rhs: Self) {
2460 self.bytes -= rhs.bytes;
2461 self.chars -= rhs.chars;
2462 }
2463 }
2464
2465 type EncodeFunc = fn(PyStrRef) -> PendingWrite;
2467 const fn textio_encode_utf8(s: PyStrRef) -> PendingWrite {
2468 PendingWrite::Utf8(s)
2469 }
2470
2471 #[derive(Debug)]
2472 struct TextIOData {
2473 buffer: PyObjectRef,
2474 encoder: Option<(PyObjectRef, Option<EncodeFunc>)>,
2475 decoder: Option<PyObjectRef>,
2476 encoding: PyUtf8StrRef,
2477 errors: PyUtf8StrRef,
2478 newline: Newlines,
2479 line_buffering: bool,
2480 write_through: bool,
2481 chunk_size: usize,
2482 seekable: bool,
2483 has_read1: bool,
2484 pending: PendingWrites,
2486 telling: bool,
2487 snapshot: Option<(i32, PyBytesRef)>,
2488 decoded_chars: Option<PyStrRef>,
2489 decoded_chars_used: Utf8size,
2491 b2cratio: f64,
2492 }
2493
2494 #[derive(Debug, Default)]
2495 struct PendingWrites {
2496 num_bytes: usize,
2497 data: PendingWritesData,
2498 }
2499
2500 #[derive(Debug, Default)]
2501 enum PendingWritesData {
2502 #[default]
2503 None,
2504 One(PendingWrite),
2505 Many(Vec<PendingWrite>),
2506 }
2507
2508 #[derive(Debug)]
2509 enum PendingWrite {
2510 Utf8(PyStrRef),
2511 Bytes(PyBytesRef),
2512 }
2513
2514 impl PendingWrite {
2515 fn as_bytes(&self) -> &[u8] {
2516 match self {
2517 Self::Utf8(s) => s.as_bytes(),
2518 Self::Bytes(b) => b.as_bytes(),
2519 }
2520 }
2521 }
2522
2523 impl PendingWrites {
2524 fn push(&mut self, write: PendingWrite) {
2525 self.num_bytes += write.as_bytes().len();
2526 self.data = match core::mem::take(&mut self.data) {
2527 PendingWritesData::None => PendingWritesData::One(write),
2528 PendingWritesData::One(write1) => PendingWritesData::Many(vec![write1, write]),
2529 PendingWritesData::Many(mut v) => {
2530 v.push(write);
2531 PendingWritesData::Many(v)
2532 }
2533 }
2534 }
2535 fn take(&mut self, vm: &VirtualMachine) -> PyBytesRef {
2536 let Self { num_bytes, data } = core::mem::take(self);
2537 if let PendingWritesData::One(PendingWrite::Bytes(b)) = data {
2538 return b;
2539 }
2540 let writes_iter = match data {
2541 PendingWritesData::None => itertools::Either::Left(vec![].into_iter()),
2542 PendingWritesData::One(write) => itertools::Either::Right(core::iter::once(write)),
2543 PendingWritesData::Many(writes) => itertools::Either::Left(writes.into_iter()),
2544 };
2545 let mut buf = Vec::with_capacity(num_bytes);
2546 writes_iter.for_each(|chunk| buf.extend_from_slice(chunk.as_bytes()));
2547 PyBytes::from(buf).into_ref(&vm.ctx)
2548 }
2549 }
2550
2551 #[derive(Default, Debug)]
2552 struct TextIOCookie {
2553 start_pos: Offset,
2554 dec_flags: i32,
2555 bytes_to_feed: i32,
2556 chars_to_skip: i32,
2557 need_eof: bool,
2558 bytes_to_skip: i32,
2560 }
2561
2562 impl TextIOCookie {
2563 const START_POS_OFF: usize = 0;
2564 const DEC_FLAGS_OFF: usize = Self::START_POS_OFF + core::mem::size_of::<Offset>();
2565 const BYTES_TO_FEED_OFF: usize = Self::DEC_FLAGS_OFF + 4;
2566 const CHARS_TO_SKIP_OFF: usize = Self::BYTES_TO_FEED_OFF + 4;
2567 const NEED_EOF_OFF: usize = Self::CHARS_TO_SKIP_OFF + 4;
2568 const BYTES_TO_SKIP_OFF: usize = Self::NEED_EOF_OFF + 1;
2569 const BYTE_LEN: usize = Self::BYTES_TO_SKIP_OFF + 4;
2570
2571 fn parse(cookie: &BigInt) -> Option<Self> {
2572 let (_, mut buf) = cookie.to_bytes_le();
2573 if buf.len() > Self::BYTE_LEN {
2574 return None;
2575 }
2576 buf.resize(Self::BYTE_LEN, 0);
2577 let buf: &[u8; Self::BYTE_LEN] = buf.as_array()?;
2578 macro_rules! get_field {
2579 ($t:ty, $off:ident) => {
2580 <$t>::from_ne_bytes(*buf[Self::$off..].first_chunk().unwrap())
2581 };
2582 }
2583 Some(Self {
2584 start_pos: get_field!(Offset, START_POS_OFF),
2585 dec_flags: get_field!(i32, DEC_FLAGS_OFF),
2586 bytes_to_feed: get_field!(i32, BYTES_TO_FEED_OFF),
2587 chars_to_skip: get_field!(i32, CHARS_TO_SKIP_OFF),
2588 need_eof: get_field!(u8, NEED_EOF_OFF) != 0,
2589 bytes_to_skip: get_field!(i32, BYTES_TO_SKIP_OFF),
2590 })
2591 }
2592
2593 fn build(&self) -> BigInt {
2594 let mut buf = [0; Self::BYTE_LEN];
2595 macro_rules! set_field {
2596 ($field:expr, $off:ident) => {{
2597 let field = $field;
2598 buf[Self::$off..][..core::mem::size_of_val(&field)]
2599 .copy_from_slice(&field.to_ne_bytes())
2600 }};
2601 }
2602 set_field!(self.start_pos, START_POS_OFF);
2603 set_field!(self.dec_flags, DEC_FLAGS_OFF);
2604 set_field!(self.bytes_to_feed, BYTES_TO_FEED_OFF);
2605 set_field!(self.chars_to_skip, CHARS_TO_SKIP_OFF);
2606 set_field!(self.need_eof as u8, NEED_EOF_OFF);
2607 set_field!(self.bytes_to_skip, BYTES_TO_SKIP_OFF);
2608 BigInt::from_signed_bytes_le(&buf)
2609 }
2610
2611 fn set_decoder_state(&self, decoder: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
2612 if self.start_pos == 0 && self.dec_flags == 0 {
2613 vm.call_method(decoder, "reset", ())?;
2614 } else {
2615 vm.call_method(
2616 decoder,
2617 "setstate",
2618 ((vm.ctx.new_bytes(vec![]), self.dec_flags),),
2619 )?;
2620 }
2621 Ok(())
2622 }
2623
2624 const fn num_to_skip(&self) -> Utf8size {
2625 Utf8size {
2626 bytes: self.bytes_to_skip as usize,
2627 chars: self.chars_to_skip as usize,
2628 }
2629 }
2630
2631 const fn set_num_to_skip(&mut self, num: Utf8size) {
2632 self.bytes_to_skip = num.bytes as i32;
2633 self.chars_to_skip = num.chars as i32;
2634 }
2635 }
2636
2637 #[pyclass(module = "_io", name, no_attr)]
2638 #[derive(Debug, PyPayload)]
2639 struct StatelessIncrementalEncoder {
2640 encode: PyObjectRef,
2641 errors: Option<PyStrRef>,
2642 name: Option<PyStrRef>,
2643 }
2644
2645 #[pyclass]
2646 impl StatelessIncrementalEncoder {
2647 #[pymethod]
2648 fn encode(
2649 &self,
2650 input: PyObjectRef,
2651 _final: OptionalArg<bool>,
2652 vm: &VirtualMachine,
2653 ) -> PyResult {
2654 let mut args: Vec<PyObjectRef> = vec![input];
2655 if let Some(errors) = &self.errors {
2656 args.push(errors.to_owned().into());
2657 }
2658 let res = self.encode.call(args, vm)?;
2659 let tuple: PyTupleRef = res.try_into_value(vm)?;
2660 if tuple.len() != 2 {
2661 return Err(vm.new_type_error("encoder must return a tuple (object, integer)"));
2662 }
2663 Ok(tuple[0].clone())
2664 }
2665
2666 #[pymethod]
2667 fn reset(&self) {}
2668
2669 #[pymethod]
2670 fn setstate(&self, _state: PyObjectRef) {}
2671
2672 #[pymethod]
2673 fn getstate(&self, vm: &VirtualMachine) -> PyObjectRef {
2674 vm.ctx.new_int(0).into()
2675 }
2676
2677 #[pygetset]
2678 fn name(&self) -> Option<PyStrRef> {
2679 self.name.clone()
2680 }
2681 }
2682
2683 #[pyclass(module = "_io", name, no_attr)]
2684 #[derive(Debug, PyPayload)]
2685 struct StatelessIncrementalDecoder {
2686 decode: PyObjectRef,
2687 errors: Option<PyStrRef>,
2688 }
2689
2690 #[pyclass]
2691 impl StatelessIncrementalDecoder {
2692 #[pymethod]
2693 fn decode(
2694 &self,
2695 input: PyObjectRef,
2696 _final: OptionalArg<bool>,
2697 vm: &VirtualMachine,
2698 ) -> PyResult {
2699 let mut args: Vec<PyObjectRef> = vec![input];
2700 if let Some(errors) = &self.errors {
2701 args.push(errors.to_owned().into());
2702 }
2703 let res = self.decode.call(args, vm)?;
2704 let tuple: PyTupleRef = res.try_into_value(vm)?;
2705 if tuple.len() != 2 {
2706 return Err(vm.new_type_error("decoder must return a tuple (object, integer)"));
2707 }
2708 Ok(tuple[0].clone())
2709 }
2710
2711 #[pymethod]
2712 fn getstate(&self, vm: &VirtualMachine) -> (PyBytesRef, u64) {
2713 (vm.ctx.empty_bytes.to_owned(), 0)
2714 }
2715
2716 #[pymethod]
2717 fn setstate(&self, _state: PyTupleRef, _vm: &VirtualMachine) {}
2718
2719 #[pymethod]
2720 fn reset(&self) {}
2721 }
2722
2723 #[pyattr]
2724 #[pyclass(name = "TextIOWrapper", base = _TextIOBase)]
2725 #[derive(Debug, Default)]
2726 struct TextIOWrapper {
2727 _base: _TextIOBase,
2728 data: PyThreadMutex<Option<TextIOData>>,
2729 finalizing: AtomicBool,
2730 }
2731
2732 impl DefaultConstructor for TextIOWrapper {}
2733
2734 impl Initializer for TextIOWrapper {
2735 type Args = (PyObjectRef, TextIOWrapperArgs);
2736
2737 fn init(
2738 zelf: PyRef<Self>,
2739 (buffer, args): Self::Args,
2740 vm: &VirtualMachine,
2741 ) -> PyResult<()> {
2742 let mut data = zelf.lock_opt(vm)?;
2743 *data = None;
2744
2745 let encoding = Self::resolve_encoding(args.encoding, vm)?;
2746
2747 let errors = args.errors.unwrap_or_else(|| vm.ctx.new_utf8_str("strict"));
2748 Self::validate_errors(&errors, vm)?;
2749
2750 let has_read1 = vm.get_attribute_opt(buffer.clone(), "read1")?.is_some();
2751 let seekable = vm.call_method(&buffer, "seekable", ())?.try_to_bool(vm)?;
2752
2753 let newline = match args.newline {
2754 OptionalArg::Missing | OptionalArg::Present(None) => Newlines::default(),
2755 OptionalArg::Present(Some(newline)) => newline,
2756 };
2757 let (encoder, decoder) =
2758 Self::find_coder(&buffer, encoding.as_str(), &errors, newline, vm)?;
2759 if let Some((encoder, _)) = &encoder {
2760 Self::adjust_encoder_state_for_bom(encoder, encoding.as_str(), &buffer, vm)?;
2761 }
2762
2763 let line_buffering = match args.line_buffering {
2764 OptionalArg::Missing => false,
2765 OptionalArg::Present(None) => false,
2766 OptionalArg::Present(Some(value)) => value.try_to_bool(vm)?,
2767 };
2768 let write_through = match args.write_through {
2769 OptionalArg::Missing => false,
2770 OptionalArg::Present(None) => false,
2771 OptionalArg::Present(Some(value)) => value.try_to_bool(vm)?,
2772 };
2773
2774 *data = Some(TextIOData {
2775 buffer,
2776 encoder,
2777 decoder,
2778 encoding,
2779 errors,
2780 newline,
2781 line_buffering,
2782 write_through,
2783 chunk_size: 8192,
2784 seekable,
2785 has_read1,
2786
2787 pending: PendingWrites::default(),
2788 telling: seekable,
2789 snapshot: None,
2790 decoded_chars: None,
2791 decoded_chars_used: Utf8size::default(),
2792 b2cratio: 0.0,
2793 });
2794
2795 Ok(())
2796 }
2797
2798 fn slot_init(zelf: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
2799 let zelf_ref: PyRef<Self> = zelf.try_into_value(vm)?;
2800 {
2801 let mut data = zelf_ref.lock_opt(vm)?;
2802 *data = None;
2803 }
2804 let (buffer, text_args): (PyObjectRef, TextIOWrapperArgs) = args.bind(vm)?;
2805 Self::init(zelf_ref, (buffer, text_args), vm)
2806 }
2807 }
2808
2809 impl TextIOWrapper {
2810 fn lock_opt(
2811 &self,
2812 vm: &VirtualMachine,
2813 ) -> PyResult<PyThreadMutexGuard<'_, Option<TextIOData>>> {
2814 self.data
2815 .lock_wrapped(|do_lock| vm.allow_threads(do_lock))
2816 .ok_or_else(|| vm.new_runtime_error("reentrant call inside textio"))
2817 }
2818
2819 fn lock(&self, vm: &VirtualMachine) -> PyResult<PyMappedThreadMutexGuard<'_, TextIOData>> {
2820 let lock = self.lock_opt(vm)?;
2821 PyThreadMutexGuard::try_map(lock, |x| x.as_mut())
2822 .map_err(|_| vm.new_value_error("I/O operation on uninitialized object"))
2823 }
2824
2825 fn validate_errors(errors: &PyRef<PyUtf8Str>, vm: &VirtualMachine) -> PyResult<()> {
2826 if errors.as_str().contains('\0') {
2827 return Err(cstring_error(vm));
2828 }
2829 vm.state
2830 .codec_registry
2831 .lookup_error(errors.as_str(), vm)
2832 .map(drop)
2833 }
2834
2835 fn bool_from_index(value: PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
2836 let int = value.try_index(vm)?;
2837 let value: i32 = int.try_to_primitive(vm)?;
2838 Ok(value != 0)
2839 }
2840
2841 fn resolve_encoding(
2842 encoding: Option<PyUtf8StrRef>,
2843 vm: &VirtualMachine,
2844 ) -> PyResult<PyUtf8StrRef> {
2845 let encoding = match encoding {
2850 None if vm.state.config.settings.utf8_mode > 0 => {
2851 identifier_utf8!(vm, utf_8).to_owned()
2852 }
2853 Some(enc) if enc.as_str() == "locale" => match vm.import("locale", 0) {
2854 Ok(locale) => locale
2855 .get_attr("getencoding", vm)?
2856 .call((), vm)?
2857 .try_into_value(vm)?,
2858 Err(err)
2859 if err.fast_isinstance(vm.ctx.exceptions.import_error)
2860 || err.fast_isinstance(vm.ctx.exceptions.module_not_found_error) =>
2861 {
2862 identifier_utf8!(vm, utf_8).to_owned()
2863 }
2864 Err(err) => return Err(err),
2865 },
2866 Some(enc) => {
2867 if enc.as_str().contains('\0') {
2868 return Err(cstring_error(vm));
2869 }
2870 enc
2871 }
2872 _ => match vm.import("locale", 0) {
2873 Ok(locale) => locale
2874 .get_attr("getencoding", vm)?
2875 .call((), vm)?
2876 .try_into_value(vm)?,
2877 Err(err)
2878 if err.fast_isinstance(vm.ctx.exceptions.import_error)
2879 || err.fast_isinstance(vm.ctx.exceptions.module_not_found_error) =>
2880 {
2881 identifier_utf8!(vm, utf_8).to_owned()
2882 }
2883 Err(err) => return Err(err),
2884 },
2885 };
2886 if encoding.as_str().contains('\0') {
2887 return Err(cstring_error(vm));
2888 }
2889 Ok(encoding)
2890 }
2891
2892 fn adjust_encoder_state_for_bom(
2893 encoder: &PyObjectRef,
2894 encoding: &str,
2895 buffer: &PyObject,
2896 vm: &VirtualMachine,
2897 ) -> PyResult<()> {
2898 let needs_bom = matches!(encoding, "utf-8-sig" | "utf-16" | "utf-32");
2899 if !needs_bom {
2900 return Ok(());
2901 }
2902 let seekable = vm.call_method(buffer, "seekable", ())?.try_to_bool(vm)?;
2903 if !seekable {
2904 return Ok(());
2905 }
2906 let pos = vm.call_method(buffer, "tell", ())?;
2907 if vm.bool_eq(&pos, vm.ctx.new_int(0).as_ref())? {
2908 return Ok(());
2909 }
2910 if let Err(err) = vm.call_method(encoder, "setstate", (0,))
2911 && !err.fast_isinstance(vm.ctx.exceptions.attribute_error)
2912 {
2913 return Err(err);
2914 }
2915 Ok(())
2916 }
2917
2918 #[allow(clippy::type_complexity)]
2919 fn find_coder(
2920 buffer: &PyObject,
2921 encoding: &str,
2922 errors: &Py<PyUtf8Str>,
2923 newline: Newlines,
2924 vm: &VirtualMachine,
2925 ) -> PyResult<(
2926 Option<(PyObjectRef, Option<EncodeFunc>)>,
2927 Option<PyObjectRef>,
2928 )> {
2929 let codec = vm.state.codec_registry.lookup(encoding, vm)?;
2930 if !codec.is_text_codec(vm)? {
2931 return Err(vm.new_lookup_error(format!(
2932 "'{encoding}' is not a text encoding; use codecs.open() to handle arbitrary codecs"
2933 )));
2934 }
2935 let errors = errors.to_owned().into_wtf8();
2936
2937 let encoder = if vm.call_method(buffer, "writable", ())?.try_to_bool(vm)? {
2938 let incremental_encoder =
2939 match codec.get_incremental_encoder(Some(errors.clone()), vm) {
2940 Ok(encoder) => encoder,
2941 Err(err)
2942 if err.fast_isinstance(vm.ctx.exceptions.type_error)
2943 || err.fast_isinstance(vm.ctx.exceptions.attribute_error) =>
2944 {
2945 let name = vm
2946 .get_attribute_opt(codec.as_tuple().to_owned().into(), "name")?
2947 .and_then(|obj| obj.downcast::<PyStr>().ok());
2948 StatelessIncrementalEncoder {
2949 encode: codec.get_encode_func().to_owned(),
2950 errors: Some(errors.clone()),
2951 name,
2952 }
2953 .into_ref(&vm.ctx)
2954 .into()
2955 }
2956 Err(err) => return Err(err),
2957 };
2958 let encoding_name = vm.get_attribute_opt(incremental_encoder.clone(), "name")?;
2959 let encode_func = encoding_name.and_then(|name| {
2960 let name = name.downcast_ref::<PyStr>()?;
2961 match name.to_str()? {
2962 "utf-8" => Some(textio_encode_utf8 as EncodeFunc),
2963 _ => None,
2964 }
2965 });
2966 Some((incremental_encoder, encode_func))
2967 } else {
2968 None
2969 };
2970
2971 let decoder = if vm.call_method(buffer, "readable", ())?.try_to_bool(vm)? {
2972 let decoder = match codec.get_incremental_decoder(Some(errors.clone()), vm) {
2973 Ok(decoder) => decoder,
2974 Err(err)
2975 if err.fast_isinstance(vm.ctx.exceptions.type_error)
2976 || err.fast_isinstance(vm.ctx.exceptions.attribute_error) =>
2977 {
2978 StatelessIncrementalDecoder {
2979 decode: codec.get_decode_func().to_owned(),
2980 errors: Some(errors),
2981 }
2982 .into_ref(&vm.ctx)
2983 .into()
2984 }
2985 Err(err) => return Err(err),
2986 };
2987 if let Newlines::Universal | Newlines::Passthrough = newline {
2988 let args = IncrementalNewlineDecoderArgs {
2989 decoder,
2990 translate: matches!(newline, Newlines::Universal),
2991 errors: None,
2992 };
2993 Some(IncrementalNewlineDecoder::construct_and_init(args, vm)?.into())
2994 } else {
2995 Some(decoder)
2996 }
2997 } else {
2998 None
2999 };
3000 Ok((encoder, decoder))
3001 }
3002 }
3003
3004 #[inline]
3005 fn flush_inner(textio: &mut TextIOData, vm: &VirtualMachine) -> PyResult {
3006 textio.check_closed(vm)?;
3007 textio.telling = textio.seekable;
3008 textio.write_pending(vm)?;
3009 vm.call_method(&textio.buffer, "flush", ())
3010 }
3011
3012 #[pyclass(
3013 with(
3014 Constructor,
3015 Initializer,
3016 Destructor,
3017 Iterable,
3018 IterNext,
3019 Representable
3020 ),
3021 flags(BASETYPE, HAS_WEAKREF)
3022 )]
3023 impl TextIOWrapper {
3024 #[pymethod]
3025 fn reconfigure(&self, args: TextIOWrapperArgs, vm: &VirtualMachine) -> PyResult<()> {
3026 let mut data = self.lock(vm)?;
3027 data.check_closed(vm)?;
3028
3029 let mut encoding = data.encoding.clone();
3030 let mut errors = data.errors.clone();
3031 let mut newline = data.newline;
3032 let mut encoding_changed = false;
3033 let mut errors_changed = false;
3034 let mut newline_changed = false;
3035 let mut line_buffering = None;
3036 let mut write_through = None;
3037
3038 if let Some(enc) = args.encoding {
3039 if enc.as_str().contains('\0') && enc.as_str().starts_with("locale") {
3040 return Err(vm.new_lookup_error(format!("unknown encoding: {enc}")));
3041 }
3042 let resolved = Self::resolve_encoding(Some(enc), vm)?;
3043 encoding_changed = resolved.as_str() != encoding.as_str();
3044 encoding = resolved;
3045 }
3046
3047 if let Some(errs) = args.errors {
3048 Self::validate_errors(&errs, vm)?;
3049 errors_changed = errs.as_str() != errors.as_str();
3050 errors = errs;
3051 } else if encoding_changed {
3052 errors = identifier_utf8!(vm, strict).to_owned();
3053 errors_changed = true;
3054 }
3055
3056 if let OptionalArg::Present(nl) = args.newline {
3057 let nl = nl.unwrap_or_default();
3058 newline_changed = nl != newline;
3059 newline = nl;
3060 }
3061
3062 if let OptionalArg::Present(Some(value)) = args.line_buffering {
3063 line_buffering = Some(Self::bool_from_index(value, vm)?);
3064 }
3065 if let OptionalArg::Present(Some(value)) = args.write_through {
3066 write_through = Some(Self::bool_from_index(value, vm)?);
3067 }
3068
3069 if (encoding_changed || newline_changed)
3070 && data.decoder.is_some()
3071 && (data.decoded_chars.is_some()
3072 || data.snapshot.is_some()
3073 || data.decoded_chars_used.chars != 0)
3074 {
3075 return Err(new_unsupported_operation(
3076 vm,
3077 "cannot reconfigure encoding or newline after reading from the stream"
3078 .to_owned(),
3079 ));
3080 }
3081
3082 if data.pending.num_bytes > 0 {
3083 data.write_pending(vm)?;
3084 }
3085 vm.call_method(&data.buffer, "flush", ())?;
3086
3087 if encoding_changed || errors_changed || newline_changed {
3088 if data.pending.num_bytes > 0 {
3089 data.write_pending(vm)?;
3090 }
3091 let (encoder, decoder) =
3092 Self::find_coder(&data.buffer, encoding.as_str(), &errors, newline, vm)?;
3093 data.encoding = encoding;
3094 data.errors = errors;
3095 data.newline = newline;
3096 data.encoder = encoder;
3097 data.decoder = decoder;
3098 data.set_decoded_chars(None);
3099 data.snapshot = None;
3100 data.decoded_chars_used = Utf8size::default();
3101 if let Some((encoder, _)) = &data.encoder {
3102 Self::adjust_encoder_state_for_bom(
3103 encoder,
3104 data.encoding.as_str(),
3105 &data.buffer,
3106 vm,
3107 )?;
3108 }
3109 }
3110
3111 if let Some(line_buffering) = line_buffering {
3112 data.line_buffering = line_buffering;
3113 }
3114 if let Some(write_through) = write_through {
3115 data.write_through = write_through;
3116 }
3117 Ok(())
3118 }
3119
3120 #[pymethod]
3121 fn detach(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
3122 let mut textio = zelf.lock(vm)?;
3123
3124 if vm.is_none(&textio.buffer) {
3126 return Err(vm.new_value_error("underlying buffer has been detached"));
3127 }
3128
3129 flush_inner(&mut textio, vm)?;
3130
3131 let buffer = textio.buffer.clone();
3132 textio.buffer = vm.ctx.none();
3133 Ok(buffer)
3134 }
3135
3136 #[pymethod]
3137 fn seekable(&self, vm: &VirtualMachine) -> PyResult {
3138 let textio = self.lock(vm)?;
3139 vm.call_method(&textio.buffer, "seekable", ())
3140 }
3141
3142 #[pymethod]
3143 fn readable(&self, vm: &VirtualMachine) -> PyResult {
3144 let textio = self.lock(vm)?;
3145 vm.call_method(&textio.buffer, "readable", ())
3146 }
3147
3148 #[pymethod]
3149 fn writable(&self, vm: &VirtualMachine) -> PyResult {
3150 let textio = self.lock(vm)?;
3151 vm.call_method(&textio.buffer, "writable", ())
3152 }
3153
3154 #[pygetset]
3155 fn line_buffering(&self, vm: &VirtualMachine) -> PyResult<bool> {
3156 Ok(self.lock(vm)?.line_buffering)
3157 }
3158
3159 #[pygetset]
3160 fn write_through(&self, vm: &VirtualMachine) -> PyResult<bool> {
3161 Ok(self.lock(vm)?.write_through)
3162 }
3163
3164 #[pygetset]
3165 fn newlines(&self, vm: &VirtualMachine) -> PyResult<Option<PyObjectRef>> {
3166 let data = self.lock(vm)?;
3167 let Some(decoder) = &data.decoder else {
3168 return Ok(None);
3169 };
3170 vm.get_attribute_opt(decoder.clone(), "newlines")
3171 }
3172
3173 #[pygetset(name = "_CHUNK_SIZE")]
3174 fn chunksize(&self, vm: &VirtualMachine) -> PyResult<usize> {
3175 Ok(self.lock(vm)?.chunk_size)
3176 }
3177
3178 #[pygetset(setter, name = "_CHUNK_SIZE")]
3179 fn set_chunksize(
3180 &self,
3181 chunk_size: PySetterValue<usize>,
3182 vm: &VirtualMachine,
3183 ) -> PyResult<()> {
3184 let mut textio = self.lock(vm)?;
3185 match chunk_size {
3186 PySetterValue::Assign(chunk_size) => textio.chunk_size = chunk_size,
3187 PySetterValue::Delete => Err(vm.new_attribute_error("cannot delete attribute"))?,
3188 };
3189 Ok(())
3193 }
3194
3195 #[pymethod]
3196 fn seek(
3197 zelf: PyRef<Self>,
3198 cookie: PyObjectRef,
3199 how: OptionalArg<i32>,
3200 vm: &VirtualMachine,
3201 ) -> PyResult {
3202 let how = how.unwrap_or(0);
3203
3204 let reset_encoder = |encoder, start_of_stream| {
3205 if start_of_stream {
3206 vm.call_method(encoder, "reset", ())
3207 } else {
3208 vm.call_method(encoder, "setstate", (0,))
3209 }
3210 };
3211
3212 let textio = zelf.lock(vm)?;
3213
3214 if !textio.seekable {
3215 return Err(new_unsupported_operation(
3216 vm,
3217 "underlying stream is not seekable".to_owned(),
3218 ));
3219 }
3220
3221 let cookie = match how {
3222 0 => cookie,
3224 1 => {
3226 if vm.bool_eq(&cookie, vm.ctx.new_int(0).as_ref())? {
3227 vm.call_method(&textio.buffer, "tell", ())?
3228 } else {
3229 return Err(new_unsupported_operation(
3230 vm,
3231 "can't do nonzero cur-relative seeks".to_owned(),
3232 ));
3233 }
3234 }
3235 2 => {
3237 if vm.bool_eq(&cookie, vm.ctx.new_int(0).as_ref())? {
3238 drop(textio);
3239 vm.call_method(zelf.as_object(), "flush", ())?;
3240 let mut textio = zelf.lock(vm)?;
3241 textio.set_decoded_chars(None);
3242 textio.snapshot = None;
3243 if let Some(decoder) = &textio.decoder {
3244 vm.call_method(decoder, "reset", ())?;
3245 }
3246 let res = vm.call_method(&textio.buffer, "seek", (0, 2))?;
3247 if let Some((encoder, _)) = &textio.encoder {
3248 let start_of_stream = vm.bool_eq(&res, vm.ctx.new_int(0).as_ref())?;
3249 reset_encoder(encoder, start_of_stream)?;
3250 }
3251 return Ok(res);
3252 } else {
3253 return Err(new_unsupported_operation(
3254 vm,
3255 "can't do nonzero end-relative seeks".to_owned(),
3256 ));
3257 }
3258 }
3259 _ => {
3260 return Err(
3261 vm.new_value_error(format!("invalid whence ({how}, should be 0, 1 or 2)"))
3262 );
3263 }
3264 };
3265 use crate::types::PyComparisonOp;
3266 if cookie.rich_compare_bool(vm.ctx.new_int(0).as_ref(), PyComparisonOp::Lt, vm)? {
3267 return Err(
3268 vm.new_value_error(format!("negative seek position {}", &cookie.repr(vm)?))
3269 );
3270 }
3271 drop(textio);
3272 vm.call_method(zelf.as_object(), "flush", ())?;
3273 let cookie_obj = crate::builtins::PyIntRef::try_from_object(vm, cookie)?;
3274 let cookie = TextIOCookie::parse(cookie_obj.as_bigint())
3275 .ok_or_else(|| vm.new_value_error("invalid cookie"))?;
3276 let mut textio = zelf.lock(vm)?;
3277 vm.call_method(&textio.buffer, "seek", (cookie.start_pos,))?;
3278 textio.set_decoded_chars(None);
3279 textio.snapshot = None;
3280 if let Some(decoder) = &textio.decoder {
3281 cookie.set_decoder_state(decoder, vm)?;
3282 }
3283 if cookie.chars_to_skip != 0 {
3284 let TextIOData {
3285 ref decoder,
3286 ref buffer,
3287 ref mut snapshot,
3288 ..
3289 } = *textio;
3290 let decoder = decoder
3291 .as_ref()
3292 .ok_or_else(|| vm.new_value_error("invalid cookie"))?;
3293 let input_chunk = vm.call_method(buffer, "read", (cookie.bytes_to_feed,))?;
3294 let input_chunk: PyBytesRef = input_chunk.downcast().map_err(|obj| {
3295 vm.new_type_error(format!(
3296 "underlying read() should have returned a bytes object, not '{}'",
3297 obj.class().name()
3298 ))
3299 })?;
3300 *snapshot = Some((cookie.dec_flags, input_chunk.clone()));
3301 let decoded = vm.call_method(decoder, "decode", (input_chunk, cookie.need_eof))?;
3302 let decoded = check_decoded(decoded, vm)?;
3303 let pos_is_valid = decoded
3304 .as_wtf8()
3305 .is_code_point_boundary(cookie.bytes_to_skip as usize);
3306 textio.set_decoded_chars(Some(decoded));
3307 if !pos_is_valid {
3308 return Err(vm.new_os_error("can't restore logical file position"));
3309 }
3310 textio.decoded_chars_used = cookie.num_to_skip();
3311 } else {
3312 textio.snapshot = Some((cookie.dec_flags, PyBytes::from(vec![]).into_ref(&vm.ctx)))
3313 }
3314 if let Some((encoder, _)) = &textio.encoder {
3315 let start_of_stream = cookie.start_pos == 0 && cookie.dec_flags == 0;
3316 reset_encoder(encoder, start_of_stream)?;
3317 }
3318 Ok(cookie_obj.into())
3319 }
3320
3321 #[pymethod]
3322 fn tell(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
3323 let mut textio = zelf.lock(vm)?;
3324 if !textio.seekable {
3325 return Err(new_unsupported_operation(
3326 vm,
3327 "underlying stream is not seekable".to_owned(),
3328 ));
3329 }
3330 if !textio.telling {
3331 return Err(vm.new_os_error("telling position disabled by next() call"));
3332 }
3333 textio.write_pending(vm)?;
3334 drop(textio);
3335 vm.call_method(zelf.as_object(), "flush", ())?;
3336 let textio = zelf.lock(vm)?;
3337 let pos = vm.call_method(&textio.buffer, "tell", ())?;
3338 let (decoder, (dec_flags, next_input)) = match (&textio.decoder, &textio.snapshot) {
3339 (Some(d), Some(s)) => (d, s),
3340 _ => return Ok(pos),
3341 };
3342 let pos = Offset::try_from_object(vm, pos)?;
3343 let mut cookie = TextIOCookie {
3344 start_pos: pos - next_input.len() as Offset,
3345 dec_flags: *dec_flags,
3346 ..Default::default()
3347 };
3348 if textio.decoded_chars_used.bytes == 0 {
3349 return Ok(cookie.build().to_pyobject(vm));
3350 }
3351 let decoder_getstate = || {
3352 let state = vm.call_method(decoder, "getstate", ())?;
3353 parse_decoder_state(state, vm)
3354 };
3355 let decoder_decode = |b: &[u8]| {
3356 let decoded = vm.call_method(decoder, "decode", (vm.ctx.new_bytes(b.to_vec()),))?;
3357 let decoded = check_decoded(decoded, vm)?;
3358 Ok(Utf8size::len_pystr(&decoded))
3359 };
3360 let saved_state = vm.call_method(decoder, "getstate", ())?;
3361 let mut num_to_skip = textio.decoded_chars_used;
3362 let mut skip_bytes = (textio.b2cratio * num_to_skip.chars as f64) as isize;
3363 let mut skip_back = 1;
3364 while skip_bytes > 0 {
3365 cookie.set_decoder_state(decoder, vm)?;
3366 let input = &next_input.as_bytes()[..skip_bytes as usize];
3367 let n_decoded = decoder_decode(input)?;
3368 if n_decoded.chars <= num_to_skip.chars {
3369 let (dec_buffer, dec_flags) = decoder_getstate()?;
3370 if dec_buffer.is_empty() {
3371 cookie.dec_flags = dec_flags;
3372 num_to_skip -= n_decoded;
3373 break;
3374 }
3375 skip_bytes -= dec_buffer.len() as isize;
3376 skip_back = 1;
3377 } else {
3378 skip_bytes -= skip_back;
3379 skip_back *= 2;
3380 }
3381 }
3382 if skip_bytes <= 0 {
3383 skip_bytes = 0;
3384 cookie.set_decoder_state(decoder, vm)?;
3385 }
3386 let skip_bytes = skip_bytes as usize;
3387
3388 cookie.start_pos += skip_bytes as Offset;
3389 cookie.set_num_to_skip(num_to_skip);
3390
3391 if num_to_skip.chars != 0 {
3392 let mut n_decoded = Utf8size::default();
3393 let mut input = next_input.as_bytes();
3394 input = &input[skip_bytes..];
3395 while !input.is_empty() {
3396 let (byte1, rest) = input.split_at(1);
3397 let n = decoder_decode(byte1)?;
3398 n_decoded += n;
3399 cookie.bytes_to_feed += 1;
3400 let (dec_buffer, dec_flags) = decoder_getstate()?;
3401 if dec_buffer.is_empty() && n_decoded.chars <= num_to_skip.chars {
3402 cookie.start_pos += cookie.bytes_to_feed as Offset;
3403 num_to_skip -= n_decoded;
3404 cookie.dec_flags = dec_flags;
3405 cookie.bytes_to_feed = 0;
3406 n_decoded = Utf8size::default();
3407 }
3408 if n_decoded.chars >= num_to_skip.chars {
3409 break;
3410 }
3411 input = rest;
3412 }
3413 if input.is_empty() {
3414 let decoded =
3415 vm.call_method(decoder, "decode", (vm.ctx.new_bytes(vec![]), true))?;
3416 let decoded = check_decoded(decoded, vm)?;
3417 let final_decoded_chars = n_decoded.chars + decoded.char_len();
3418 cookie.need_eof = true;
3419 if final_decoded_chars < num_to_skip.chars {
3420 return Err(vm.new_os_error("can't reconstruct logical file position"));
3421 }
3422 }
3423 }
3424 vm.call_method(decoder, "setstate", (saved_state,))?;
3425 cookie.set_num_to_skip(num_to_skip);
3426 Ok(cookie.build().to_pyobject(vm))
3427 }
3428
3429 #[pygetset]
3430 fn name(&self, vm: &VirtualMachine) -> PyResult {
3431 let buffer = self.lock(vm)?.buffer.clone();
3432 buffer.get_attr("name", vm)
3433 }
3434
3435 #[pygetset]
3436 fn encoding(&self, vm: &VirtualMachine) -> PyResult<PyUtf8StrRef> {
3437 Ok(self.lock(vm)?.encoding.clone())
3438 }
3439
3440 #[pygetset]
3441 fn errors(&self, vm: &VirtualMachine) -> PyResult<PyUtf8StrRef> {
3442 Ok(self.lock(vm)?.errors.clone())
3443 }
3444
3445 #[pymethod]
3446 fn fileno(&self, vm: &VirtualMachine) -> PyResult {
3447 let buffer = self.lock(vm)?.buffer.clone();
3448 vm.call_method(&buffer, "fileno", ())
3449 }
3450
3451 #[pymethod]
3452 fn read(&self, size: OptionalSize, vm: &VirtualMachine) -> PyResult<PyStrRef> {
3453 let mut textio = self.lock(vm)?;
3454 textio.check_closed(vm)?;
3455 let decoder = textio
3456 .decoder
3457 .clone()
3458 .ok_or_else(|| new_unsupported_operation(vm, "not readable".to_owned()))?;
3459
3460 textio.write_pending(vm)?;
3461
3462 let s = if let Some(mut remaining) = size.to_usize() {
3463 let mut chunks = Vec::new();
3464 let mut chunks_bytes = 0;
3465 loop {
3466 if let Some((s, char_len)) = textio.get_decoded_chars(remaining, vm) {
3467 chunks_bytes += s.byte_len();
3468 chunks.push(s);
3469 remaining = remaining.saturating_sub(char_len);
3470 }
3471 if remaining == 0 {
3472 break;
3473 }
3474 let eof = textio.read_chunk(remaining, vm)?;
3475 if eof {
3476 break;
3477 }
3478 }
3479 if chunks.is_empty() {
3480 vm.ctx.empty_str.to_owned()
3481 } else if chunks.len() == 1 {
3482 chunks.pop().unwrap()
3483 } else {
3484 let mut ret = Wtf8Buf::with_capacity(chunks_bytes);
3485 for chunk in chunks {
3486 ret.push_wtf8(chunk.as_wtf8())
3487 }
3488 PyStr::from(ret).into_ref(&vm.ctx)
3489 }
3490 } else {
3491 let bytes = vm.call_method(&textio.buffer, "read", ())?;
3492 let decoded = vm.call_method(&decoder, "decode", (bytes, true))?;
3493 let decoded = check_decoded(decoded, vm)?;
3494 let ret = textio.take_decoded_chars(Some(decoded), vm);
3495 textio.snapshot = None;
3496 ret
3497 };
3498 Ok(s)
3499 }
3500
3501 #[pymethod]
3502 fn write(&self, obj: PyStrRef, vm: &VirtualMachine) -> PyResult<usize> {
3503 let mut textio = self.lock(vm)?;
3504 textio.check_closed(vm)?;
3505
3506 let (encoder, encode_func) = textio
3507 .encoder
3508 .as_ref()
3509 .ok_or_else(|| new_unsupported_operation(vm, "not writable".to_owned()))?;
3510
3511 let char_len = obj.char_len();
3512
3513 let data = obj.as_wtf8();
3514
3515 let replace_nl = match textio.newline {
3516 Newlines::Lf => Some("\n"),
3517 Newlines::Cr => Some("\r"),
3518 Newlines::Crlf => Some("\r\n"),
3519 Newlines::Universal if cfg!(windows) => Some("\r\n"),
3520 _ => None,
3521 };
3522 let has_lf = (replace_nl.is_some() || textio.line_buffering)
3523 && data.contains_code_point('\n'.into());
3524 let flush = textio.line_buffering && (has_lf || data.contains_code_point('\r'.into()));
3525 let chunk = if let Some(replace_nl) = replace_nl {
3526 if has_lf {
3527 PyStr::from(data.replace("\n".as_ref(), replace_nl.as_ref())).into_ref(&vm.ctx)
3528 } else {
3529 obj
3530 }
3531 } else {
3532 obj
3533 };
3534 let chunk = if let Some(encode_func) = *encode_func {
3535 encode_func(chunk)
3536 } else {
3537 let b = vm.call_method(encoder, "encode", (chunk.clone(),))?;
3538 b.downcast::<PyBytes>()
3539 .map(PendingWrite::Bytes)
3540 .or_else(|obj| {
3541 if obj.is(&chunk) {
3544 Ok(PendingWrite::Utf8(chunk))
3545 } else {
3546 Err(vm.new_type_error(format!(
3547 "encoder should return a bytes object, not '{}'",
3548 obj.class().name()
3549 )))
3550 }
3551 })?
3552 };
3553 if textio.pending.num_bytes > 0
3554 && textio.pending.num_bytes + chunk.as_bytes().len() > textio.chunk_size
3555 {
3556 let buffer = textio.buffer.clone();
3557 let pending = textio.pending.take(vm);
3558 drop(textio);
3559 vm.call_method(&buffer, "write", (pending,))?;
3560 textio = self.lock(vm)?;
3561 textio.check_closed(vm)?;
3562 if textio.pending.num_bytes > 0 {
3563 let buffer = textio.buffer.clone();
3564 let pending = textio.pending.take(vm);
3565 drop(textio);
3566 vm.call_method(&buffer, "write", (pending,))?;
3567 textio = self.lock(vm)?;
3568 textio.check_closed(vm)?;
3569 }
3570 }
3571 textio.pending.push(chunk);
3572 if textio.pending.num_bytes > 0
3573 && (flush || textio.write_through || textio.pending.num_bytes >= textio.chunk_size)
3574 {
3575 let buffer = textio.buffer.clone();
3576 let pending = textio.pending.take(vm);
3577 drop(textio);
3578 vm.call_method(&buffer, "write", (pending,))?;
3579 textio = self.lock(vm)?;
3580 textio.check_closed(vm)?;
3581 }
3582 if flush {
3583 let _ = vm.call_method(&textio.buffer, "flush", ());
3584 }
3585
3586 Ok(char_len)
3587 }
3588
3589 #[pymethod]
3590 fn flush(&self, vm: &VirtualMachine) -> PyResult {
3591 let mut textio = self.lock(vm)?;
3592 flush_inner(&mut textio, vm)
3593 }
3594
3595 #[pymethod]
3596 fn truncate(
3597 zelf: PyRef<Self>,
3598 pos: OptionalArg<PyObjectRef>,
3599 vm: &VirtualMachine,
3600 ) -> PyResult {
3601 let mut textio = zelf.lock(vm)?;
3603 flush_inner(&mut textio, vm)?;
3604 let buffer = textio.buffer.clone();
3605 drop(textio);
3606
3607 let pos = match pos.into_option() {
3608 Some(p) => p,
3609 None => vm.call_method(zelf.as_object(), "tell", ())?,
3610 };
3611 vm.call_method(&buffer, "truncate", (pos,))
3612 }
3613
3614 #[pymethod]
3615 fn isatty(&self, vm: &VirtualMachine) -> PyResult {
3616 let textio = self.lock(vm)?;
3617 textio.check_closed(vm)?;
3618 vm.call_method(&textio.buffer, "isatty", ())
3619 }
3620
3621 #[pymethod]
3622 fn readline(&self, size: OptionalSize, vm: &VirtualMachine) -> PyResult<PyStrRef> {
3623 let limit = size.to_usize();
3624
3625 let mut textio = self.lock(vm)?;
3626 check_closed(&textio.buffer, vm)?;
3627
3628 textio.write_pending(vm)?;
3629
3630 #[derive(Clone)]
3631 struct SlicedStr(PyStrRef, Range<usize>);
3632
3633 impl SlicedStr {
3634 #[inline]
3635 fn byte_len(&self) -> usize {
3636 self.1.len()
3637 }
3638
3639 #[inline]
3640 fn char_len(&self) -> usize {
3641 if self.is_full_slice() {
3642 self.0.char_len()
3643 } else {
3644 self.slice().code_points().count()
3645 }
3646 }
3647
3648 #[inline]
3649 fn is_full_slice(&self) -> bool {
3650 self.1.len() >= self.0.byte_len()
3651 }
3652
3653 #[inline]
3654 fn slice(&self) -> &Wtf8 {
3655 &self.0.as_wtf8()[self.1.clone()]
3656 }
3657
3658 #[inline]
3659 fn slice_pystr(self, vm: &VirtualMachine) -> PyStrRef {
3660 if self.is_full_slice() {
3661 self.0
3662 } else {
3663 PyStr::from(self.slice()).into_ref(&vm.ctx)
3665 }
3666 }
3667
3668 fn utf8_len(&self) -> Utf8size {
3669 Utf8size {
3670 bytes: self.byte_len(),
3671 chars: self.char_len(),
3672 }
3673 }
3674 }
3675
3676 let mut start;
3677 let mut end_pos;
3678 let mut offset_to_buffer;
3679 let mut chunked = Utf8size::default();
3680 let mut remaining: Option<SlicedStr> = None;
3681 let mut chunks = Vec::new();
3682
3683 let cur_line = 'outer: loop {
3684 let decoded_chars = loop {
3685 match textio.decoded_chars.as_ref() {
3686 Some(s) if !s.is_empty() => break s,
3687 _ => {}
3688 }
3689 let eof = textio.read_chunk(0, vm)?;
3690 if eof {
3691 textio.set_decoded_chars(None);
3692 textio.snapshot = None;
3693 start = Utf8size::default();
3694 end_pos = Utf8size::default();
3695 offset_to_buffer = Utf8size::default();
3696 break 'outer None;
3697 }
3698 };
3699 let line = match remaining.take() {
3700 None => {
3701 start = textio.decoded_chars_used;
3702 offset_to_buffer = Utf8size::default();
3703 decoded_chars.clone()
3704 }
3705 Some(remaining) => {
3706 assert_eq!(textio.decoded_chars_used.bytes, 0);
3707 offset_to_buffer = remaining.utf8_len();
3708 let decoded_chars = decoded_chars.as_wtf8();
3709 let line = if remaining.is_full_slice() {
3710 let mut line = remaining.0;
3711 line.concat_in_place(decoded_chars, vm);
3712 line
3713 } else {
3714 let remaining = remaining.slice();
3715 let mut s =
3716 Wtf8Buf::with_capacity(remaining.len() + decoded_chars.len());
3717 s.push_wtf8(remaining);
3718 s.push_wtf8(decoded_chars);
3719 PyStr::from(s).into_ref(&vm.ctx)
3720 };
3721 start = Utf8size::default();
3722 line
3723 }
3724 };
3725 let line_from_start = &line.as_wtf8()[start.bytes..];
3726 let nl_res = textio.newline.find_newline(line_from_start);
3727 match nl_res {
3728 Ok(p) | Err(p) => {
3729 end_pos = start + Utf8size::len_str(&line_from_start[..p]);
3730 if let Some(limit) = limit {
3731 if chunked.chars + end_pos.chars >= limit {
3733 end_pos = start
3734 + Utf8size {
3735 chars: limit - chunked.chars,
3736 bytes: crate::common::str::codepoint_range_end(
3737 line_from_start,
3738 limit - chunked.chars,
3739 )
3740 .unwrap(),
3741 };
3742 break Some(line);
3743 }
3744 }
3745 }
3746 }
3747 if nl_res.is_ok() {
3748 break Some(line);
3749 }
3750 if end_pos.bytes > start.bytes {
3751 let chunk = SlicedStr(line.clone(), start.bytes..end_pos.bytes);
3752 chunked += chunk.utf8_len();
3753 chunks.push(chunk);
3754 }
3755 let line_len = line.byte_len();
3756 if end_pos.bytes < line_len {
3757 remaining = Some(SlicedStr(line, end_pos.bytes..line_len));
3758 }
3759 textio.set_decoded_chars(None);
3760 };
3761
3762 let cur_line = cur_line.map(|line| {
3763 textio.decoded_chars_used = end_pos - offset_to_buffer;
3764 SlicedStr(line, start.bytes..end_pos.bytes)
3765 });
3766 let mut chunked = chunked.bytes;
3768 if let Some(remaining) = remaining {
3769 chunked += remaining.byte_len();
3770 chunks.push(remaining);
3771 }
3772 let line = if !chunks.is_empty() {
3773 if let Some(cur_line) = cur_line {
3774 chunked += cur_line.byte_len();
3775 chunks.push(cur_line);
3776 }
3777 let mut s = Wtf8Buf::with_capacity(chunked);
3778 for chunk in chunks {
3779 s.push_wtf8(chunk.slice())
3780 }
3781 PyStr::from(s).into_ref(&vm.ctx)
3782 } else if let Some(cur_line) = cur_line {
3783 cur_line.slice_pystr(vm)
3784 } else {
3785 vm.ctx.empty_str.to_owned()
3786 };
3787 Ok(line)
3788 }
3789
3790 #[pymethod]
3791 fn close(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<()> {
3792 let buffer = zelf.lock(vm)?.buffer.clone();
3793 if file_closed(&buffer, vm)? {
3794 return Ok(());
3795 }
3796 if vm.is_none(&zelf.lock(vm)?.buffer) {
3801 return Ok(());
3802 }
3803 if zelf.finalizing.load(Ordering::Relaxed) {
3804 let _ = vm.call_method(&buffer, "_dealloc_warn", (zelf.as_object().to_owned(),));
3806 }
3807 let flush_res = vm.call_method(zelf.as_object(), "flush", ()).map(drop);
3808 let close_res = vm.call_method(&buffer, "close", ()).map(drop);
3809 exception_chain(flush_res, close_res)
3810 }
3811
3812 #[pygetset]
3813 fn closed(&self, vm: &VirtualMachine) -> PyResult {
3814 let buffer = self.lock(vm)?.buffer.clone();
3815 buffer.get_attr("closed", vm)
3816 }
3817
3818 #[pygetset]
3819 fn buffer(&self, vm: &VirtualMachine) -> PyResult {
3820 Ok(self.lock(vm)?.buffer.clone())
3821 }
3822
3823 #[pymethod]
3824 fn __getstate__(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
3825 Err(vm.new_type_error(format!("cannot pickle '{}' instances", zelf.class().name())))
3826 }
3827
3828 #[pymethod]
3829 fn __reduce_ex__(zelf: PyObjectRef, proto: usize, vm: &VirtualMachine) -> PyResult {
3830 if zelf.class().is(TextIOWrapper::static_type()) {
3831 return Err(
3832 vm.new_type_error(format!("cannot pickle '{}' object", zelf.class().name()))
3833 );
3834 }
3835 let _ = proto;
3836 reduce_ex_for_subclass(zelf, vm)
3837 }
3838 }
3839
3840 fn parse_decoder_state(state: PyObjectRef, vm: &VirtualMachine) -> PyResult<(PyBytesRef, i32)> {
3841 use crate::builtins::{PyTuple, int};
3842 let state_err = || vm.new_type_error("illegal decoder state");
3843 let state = state.downcast::<PyTuple>().map_err(|_| state_err())?;
3844 match state.as_slice() {
3845 [buf, flags] => {
3846 let buf = buf.clone().downcast::<PyBytes>().map_err(|obj| {
3847 vm.new_type_error(format!(
3848 "illegal decoder state: the first item should be a bytes object, not '{}'",
3849 obj.class().name()
3850 ))
3851 })?;
3852 let flags = flags.downcast_ref::<int::PyInt>().ok_or_else(state_err)?;
3853 let flags = flags.try_to_primitive(vm)?;
3854 Ok((buf, flags))
3855 }
3856 _ => Err(state_err()),
3857 }
3858 }
3859
3860 impl TextIOData {
3861 fn write_pending(&mut self, vm: &VirtualMachine) -> PyResult<()> {
3862 if self.pending.num_bytes == 0 {
3863 return Ok(());
3864 }
3865 let data = self.pending.take(vm);
3866 vm.call_method(&self.buffer, "write", (data,))?;
3867 Ok(())
3868 }
3869
3870 fn read_chunk(&mut self, size_hint: usize, vm: &VirtualMachine) -> PyResult<bool> {
3872 let decoder = self
3873 .decoder
3874 .as_ref()
3875 .ok_or_else(|| new_unsupported_operation(vm, "not readable".to_owned()))?;
3876
3877 let dec_state = if self.telling {
3878 let state = vm.call_method(decoder, "getstate", ())?;
3879 Some(parse_decoder_state(state, vm)?)
3880 } else {
3881 None
3882 };
3883
3884 let method = if self.has_read1 { "read1" } else { "read" };
3885 let size_hint = if size_hint > 0 {
3886 (self.b2cratio.max(1.0) * size_hint as f64) as usize
3887 } else {
3888 size_hint
3889 };
3890 let chunk_size = core::cmp::max(self.chunk_size, size_hint);
3891 let input_chunk = vm.call_method(&self.buffer, method, (chunk_size,))?;
3892
3893 let buf = ArgBytesLike::try_from_borrowed_object(vm, &input_chunk).map_err(|_| {
3894 vm.new_type_error(format!(
3895 "underlying {}() should have returned a bytes-like object, not '{}'",
3896 method,
3897 input_chunk.class().name()
3898 ))
3899 })?;
3900 let nbytes = buf.borrow_buf().len();
3901 let eof = nbytes == 0;
3902 let decoded = vm.call_method(decoder, "decode", (input_chunk, eof))?;
3903 let decoded = check_decoded(decoded, vm)?;
3904
3905 let char_len = decoded.char_len();
3906 self.b2cratio = if char_len > 0 {
3907 nbytes as f64 / char_len as f64
3908 } else {
3909 0.0
3910 };
3911 let eof = if char_len > 0 { false } else { eof };
3912 self.set_decoded_chars(Some(decoded));
3913
3914 if let Some((dec_buffer, dec_flags)) = dec_state {
3915 let mut next_input = dec_buffer.as_bytes().to_vec();
3917 next_input.extend_from_slice(&buf.borrow_buf());
3918 self.snapshot = Some((dec_flags, PyBytes::from(next_input).into_ref(&vm.ctx)));
3919 }
3920
3921 Ok(eof)
3922 }
3923
3924 fn check_closed(&self, vm: &VirtualMachine) -> PyResult<()> {
3925 check_closed(&self.buffer, vm)
3926 }
3927
3928 fn get_decoded_chars(
3931 &mut self,
3932 n: usize,
3933 vm: &VirtualMachine,
3934 ) -> Option<(PyStrRef, usize)> {
3935 if n == 0 {
3936 return None;
3937 }
3938 let decoded_chars = self.decoded_chars.as_ref()?;
3939 let avail = &decoded_chars.as_wtf8()[self.decoded_chars_used.bytes..];
3940 if avail.is_empty() {
3941 return None;
3942 }
3943 let avail_chars = decoded_chars.char_len() - self.decoded_chars_used.chars;
3944 let (chars, chars_used) = if n >= avail_chars {
3945 if self.decoded_chars_used.bytes == 0 {
3946 (decoded_chars.clone(), avail_chars)
3947 } else {
3948 (PyStr::from(avail).into_ref(&vm.ctx), avail_chars)
3949 }
3950 } else {
3951 let s = crate::common::str::get_codepoints(avail, 0..n);
3952 (PyStr::from(s).into_ref(&vm.ctx), n)
3953 };
3954 self.decoded_chars_used += Utf8size {
3955 bytes: chars.byte_len(),
3956 chars: chars_used,
3957 };
3958 Some((chars, chars_used))
3959 }
3960
3961 fn set_decoded_chars(&mut self, s: Option<PyStrRef>) {
3962 self.decoded_chars = s;
3963 self.decoded_chars_used = Utf8size::default();
3964 }
3965
3966 fn take_decoded_chars(
3967 &mut self,
3968 append: Option<PyStrRef>,
3969 vm: &VirtualMachine,
3970 ) -> PyStrRef {
3971 let empty_str = || vm.ctx.empty_str.to_owned();
3972 let chars_pos = core::mem::take(&mut self.decoded_chars_used).bytes;
3973 let decoded_chars = match core::mem::take(&mut self.decoded_chars) {
3974 None => return append.unwrap_or_else(empty_str),
3975 Some(s) if s.is_empty() => return append.unwrap_or_else(empty_str),
3976 Some(s) => s,
3977 };
3978 let append_len = append.as_ref().map_or(0, |s| s.byte_len());
3979 if append_len == 0 && chars_pos == 0 {
3980 return decoded_chars;
3981 }
3982 let decoded_chars_unused = &decoded_chars.as_wtf8()[chars_pos..];
3984 let mut s = Wtf8Buf::with_capacity(decoded_chars_unused.len() + append_len);
3985 s.push_wtf8(decoded_chars_unused);
3986 if let Some(append) = append {
3987 s.push_wtf8(append.as_wtf8())
3988 }
3989 PyStr::from(s).into_ref(&vm.ctx)
3990 }
3991 }
3992
3993 impl Destructor for TextIOWrapper {
3994 fn slot_del(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
3995 if let Some(wrapper) = zelf.downcast_ref::<TextIOWrapper>() {
3996 wrapper.finalizing.store(true, Ordering::Relaxed);
3997 }
3998 iobase_finalize(zelf, vm);
3999 Ok(())
4000 }
4001
4002 #[cold]
4003 fn del(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<()> {
4004 unreachable!("slot_del is implemented")
4005 }
4006 }
4007
4008 impl Representable for TextIOWrapper {
4009 #[inline]
4010 fn repr(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyRef<PyStr>> {
4011 let type_name = zelf.class().slot_name();
4012 let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) else {
4013 return Err(
4014 vm.new_runtime_error(format!("reentrant call inside {type_name}.__repr__"))
4015 );
4016 };
4017 let Some(data) = zelf.data.lock() else {
4018 return Ok(vm.ctx.new_str(Wtf8Buf::from(format!("<{type_name}>"))));
4020 };
4021 let Some(data) = data.as_ref() else {
4022 return Err(vm.new_value_error("I/O operation on uninitialized object"));
4023 };
4024
4025 let mut result = Wtf8Buf::from(format!("<{type_name}"));
4026
4027 if let Ok(Some(name)) = vm.get_attribute_opt(data.buffer.clone(), "name") {
4029 let name_repr = name.repr(vm)?;
4030 result.push_wtf8(" name=".as_ref());
4031 result.push_wtf8(name_repr.as_wtf8());
4032 }
4033
4034 let mode_obj = match vm.get_attribute_opt(zelf.as_object().to_owned(), "mode") {
4036 Ok(Some(mode)) => Some(mode),
4037 Ok(None) | Err(_) => match vm.get_attribute_opt(data.buffer.clone(), "mode") {
4038 Ok(Some(mode)) => Some(mode),
4039 _ => None,
4040 },
4041 };
4042 if let Some(mode) = mode_obj {
4043 let mode_repr = mode.repr(vm)?;
4044 result.push_wtf8(" mode=".as_ref());
4045 result.push_wtf8(mode_repr.as_wtf8());
4046 }
4047
4048 result.push_wtf8(" encoding='".as_ref());
4050 result.push_wtf8(data.encoding.as_str().as_ref());
4051 result.push_wtf8("'>".as_ref());
4052
4053 Ok(vm.ctx.new_str(result))
4054 }
4055
4056 fn repr_str(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
4057 unreachable!("repr() is overridden directly")
4058 }
4059 }
4060
4061 impl Iterable for TextIOWrapper {
4062 fn slot_iter(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
4063 check_closed(&zelf, vm)?;
4064 Ok(zelf)
4065 }
4066
4067 fn iter(_zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyResult {
4068 unreachable!("slot_iter is implemented")
4069 }
4070 }
4071
4072 impl IterNext for TextIOWrapper {
4073 fn slot_iternext(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
4074 let textio_ref: PyRef<TextIOWrapper> =
4076 zelf.downcast_ref::<TextIOWrapper>().unwrap().to_owned();
4077 {
4078 let mut textio = textio_ref.lock(vm)?;
4079 textio.telling = false;
4080 }
4081
4082 let line = vm.call_method(zelf, "readline", ())?;
4083
4084 if !line.clone().try_to_bool(vm)? {
4085 let mut textio = textio_ref.lock(vm)?;
4087 textio.snapshot = None;
4088 textio.telling = textio.seekable;
4089 Ok(PyIterReturn::StopIteration(None))
4090 } else {
4091 Ok(PyIterReturn::Return(line))
4092 }
4093 }
4094
4095 fn next(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<PyIterReturn> {
4096 unreachable!("slot_iternext is implemented")
4097 }
4098 }
4099
4100 #[pyattr]
4101 #[pyclass(name)]
4102 #[derive(Debug, PyPayload, Default)]
4103 struct IncrementalNewlineDecoder {
4104 data: PyThreadMutex<Option<IncrementalNewlineDecoderData>>,
4106 }
4107
4108 #[derive(Debug)]
4109 struct IncrementalNewlineDecoderData {
4110 decoder: PyObjectRef,
4111 pendingcr: bool,
4114 translate: bool,
4115 seennl: SeenNewline,
4116 }
4117
4118 bitflags! {
4119 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
4120 struct SeenNewline: u8 {
4121 const LF = 1;
4122 const CR = 2;
4123 const CRLF = 4;
4124 }
4125 }
4126
4127 impl DefaultConstructor for IncrementalNewlineDecoder {}
4128
4129 #[derive(FromArgs)]
4130 struct IncrementalNewlineDecoderArgs {
4131 #[pyarg(any)]
4132 decoder: PyObjectRef,
4133 #[pyarg(any)]
4134 translate: bool,
4135 #[pyarg(any, default)]
4136 errors: Option<PyObjectRef>,
4137 }
4138
4139 impl Initializer for IncrementalNewlineDecoder {
4140 type Args = IncrementalNewlineDecoderArgs;
4141 fn init(zelf: PyRef<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> {
4142 let _ = args.errors;
4143 let mut data = zelf.lock_opt(vm)?;
4144 *data = Some(IncrementalNewlineDecoderData {
4145 decoder: args.decoder,
4146 translate: args.translate,
4147 pendingcr: false,
4148 seennl: SeenNewline::empty(),
4149 });
4150 Ok(())
4151 }
4152 }
4153
4154 #[pyclass(with(Constructor, Initializer))]
4155 impl IncrementalNewlineDecoder {
4156 fn lock_opt(
4157 &self,
4158 vm: &VirtualMachine,
4159 ) -> PyResult<PyThreadMutexGuard<'_, Option<IncrementalNewlineDecoderData>>> {
4160 self.data
4161 .lock_wrapped(|do_lock| vm.allow_threads(do_lock))
4162 .ok_or_else(|| vm.new_runtime_error("reentrant call inside nldecoder"))
4163 }
4164
4165 fn lock(
4166 &self,
4167 vm: &VirtualMachine,
4168 ) -> PyResult<PyMappedThreadMutexGuard<'_, IncrementalNewlineDecoderData>> {
4169 let lock = self.lock_opt(vm)?;
4170 PyThreadMutexGuard::try_map(lock, |x| x.as_mut())
4171 .map_err(|_| vm.new_value_error("I/O operation on uninitialized nldecoder"))
4172 }
4173
4174 #[pymethod]
4175 fn decode(&self, args: NewlineDecodeArgs, vm: &VirtualMachine) -> PyResult<PyStrRef> {
4176 self.lock(vm)?.decode(args.input, args.r#final, vm)
4177 }
4178
4179 #[pymethod]
4180 fn getstate(&self, vm: &VirtualMachine) -> PyResult<(PyObjectRef, u64)> {
4181 let data = self.lock(vm)?;
4182 let (buffer, flag) = if vm.is_none(&data.decoder) {
4183 (vm.ctx.new_bytes(vec![]).into(), 0)
4184 } else {
4185 vm.call_method(&data.decoder, "getstate", ())?
4186 .try_to_ref::<PyTuple>(vm)?
4187 .extract_tuple::<(PyObjectRef, u64)>(vm)?
4188 };
4189 let flag = (flag << 1) | (data.pendingcr as u64);
4190 Ok((buffer, flag))
4191 }
4192
4193 #[pymethod]
4194 fn setstate(&self, state: PyTupleRef, vm: &VirtualMachine) -> PyResult<()> {
4195 let mut data = self.lock(vm)?;
4196 let (buffer, flag) = state.extract_tuple::<(PyObjectRef, u64)>(vm)?;
4197 data.pendingcr = flag & 1 != 0;
4198 if !vm.is_none(&data.decoder) {
4199 vm.call_method(&data.decoder, "setstate", ((buffer, flag >> 1),))?;
4200 }
4201 Ok(())
4202 }
4203
4204 #[pymethod]
4205 fn reset(&self, vm: &VirtualMachine) -> PyResult<()> {
4206 let mut data = self.lock(vm)?;
4207 data.seennl = SeenNewline::empty();
4208 data.pendingcr = false;
4209 if !vm.is_none(&data.decoder) {
4210 vm.call_method(&data.decoder, "reset", ())?;
4211 }
4212 Ok(())
4213 }
4214
4215 #[pygetset]
4216 fn newlines(&self, vm: &VirtualMachine) -> PyResult {
4217 let data = self.lock(vm)?;
4218 Ok(match data.seennl.bits() {
4219 1 => "\n".to_pyobject(vm),
4220 2 => "\r".to_pyobject(vm),
4221 3 => ("\r", "\n").to_pyobject(vm),
4222 4 => "\r\n".to_pyobject(vm),
4223 5 => ("\n", "\r\n").to_pyobject(vm),
4224 6 => ("\r", "\r\n").to_pyobject(vm),
4225 7 => ("\r", "\n", "\r\n").to_pyobject(vm),
4226 _ => vm.ctx.none(),
4227 })
4228 }
4229 }
4230
4231 #[derive(FromArgs)]
4232 struct NewlineDecodeArgs {
4233 #[pyarg(any)]
4234 input: PyObjectRef,
4235 #[pyarg(any, default)]
4236 r#final: bool,
4237 }
4238
4239 impl IncrementalNewlineDecoderData {
4240 fn decode(
4241 &mut self,
4242 input: PyObjectRef,
4243 final_: bool,
4244 vm: &VirtualMachine,
4245 ) -> PyResult<PyStrRef> {
4246 let output = if vm.is_none(&self.decoder) {
4247 input
4248 } else {
4249 vm.call_method(&self.decoder, "decode", (input, final_))?
4250 };
4251 let orig_output: PyStrRef = output.try_into_value(vm)?;
4252 let mut output = Cow::Borrowed(orig_output.as_wtf8());
4254 if self.pendingcr && (final_ || !output.is_empty()) {
4255 output.to_mut().insert(0, '\r'.into());
4256 self.pendingcr = false;
4257 }
4258 if !final_ && let Some(s) = output.strip_suffix("\r") {
4259 output = Cow::Owned(s.to_owned());
4260 self.pendingcr = true;
4261 }
4262
4263 if output.is_empty() {
4264 return Ok(vm.ctx.empty_str.to_owned());
4265 }
4266
4267 if (self.seennl == SeenNewline::LF || self.seennl.is_empty())
4268 && !output.contains_code_point('\r'.into())
4269 {
4270 if self.seennl.is_empty() && output.contains_code_point('\n'.into()) {
4271 self.seennl.insert(SeenNewline::LF);
4272 }
4273 } else if !self.translate {
4274 let output = output.as_bytes();
4275 let mut matches = memchr::memchr2_iter(b'\r', b'\n', output);
4276 while !self.seennl.is_all() {
4277 let Some(i) = matches.next() else { break };
4278 match output[i] {
4279 b'\n' => self.seennl.insert(SeenNewline::LF),
4280 _ if output.get(i + 1) == Some(&b'\n') => {
4282 matches.next();
4283 self.seennl.insert(SeenNewline::CRLF);
4284 }
4285 _ => self.seennl.insert(SeenNewline::CR),
4286 }
4287 }
4288 } else {
4289 let bytes = output.as_bytes();
4290 let mut matches = memchr::memchr2_iter(b'\r', b'\n', bytes);
4291 let mut new_string = Wtf8Buf::with_capacity(output.len());
4292 let mut last_modification_index = 0;
4293 while let Some(cr_index) = matches.next() {
4294 if bytes[cr_index] == b'\r' {
4295 let mut next_chunk_index = cr_index + 1;
4297 if bytes.get(cr_index + 1) == Some(&b'\n') {
4298 matches.next();
4299 self.seennl.insert(SeenNewline::CRLF);
4300 next_chunk_index += 1;
4302 } else {
4303 self.seennl.insert(SeenNewline::CR);
4304 }
4305 new_string.push_wtf8(&output[last_modification_index..cr_index]);
4306 new_string.push_char('\n');
4307 last_modification_index = next_chunk_index;
4308 } else {
4309 self.seennl.insert(SeenNewline::LF);
4310 }
4311 }
4312 new_string.push_wtf8(&output[last_modification_index..]);
4313 output = Cow::Owned(new_string);
4314 }
4315
4316 Ok(match output {
4317 Cow::Borrowed(_) => orig_output,
4318 Cow::Owned(s) => vm.ctx.new_str(s),
4319 })
4320 }
4321 }
4322
4323 #[pyattr]
4324 #[pyclass(name = "StringIO", base = _TextIOBase)]
4325 #[derive(Debug)]
4326 struct StringIO {
4327 _base: _TextIOBase,
4328 buffer: PyRwLock<BufferedIO>,
4329 closed: AtomicCell<bool>,
4330 }
4331
4332 #[derive(FromArgs)]
4333 struct StringIONewArgs {
4334 #[pyarg(positional, optional)]
4335 object: OptionalOption<PyStrRef>,
4336
4337 #[pyarg(any, default)]
4339 #[allow(dead_code)]
4340 newline: Newlines,
4341 }
4342
4343 impl Constructor for StringIO {
4344 type Args = FuncArgs;
4345
4346 fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
4347 Ok(Self {
4348 _base: Default::default(),
4349 buffer: PyRwLock::new(BufferedIO::new(Cursor::new(Vec::new()))),
4350 closed: AtomicCell::new(false),
4351 })
4352 }
4353 }
4354
4355 impl Initializer for StringIO {
4356 type Args = StringIONewArgs;
4357
4358 #[allow(unused_variables)]
4359 fn init(
4360 zelf: PyRef<Self>,
4361 Self::Args { object, newline }: Self::Args,
4362 _vm: &VirtualMachine,
4363 ) -> PyResult<()> {
4364 let raw_bytes = object
4365 .flatten()
4366 .map_or_else(Vec::new, |v| v.as_bytes().to_vec());
4367 *zelf.buffer.write() = BufferedIO::new(Cursor::new(raw_bytes));
4368 Ok(())
4369 }
4370 }
4371
4372 impl StringIO {
4373 fn buffer(&self, vm: &VirtualMachine) -> PyResult<PyRwLockWriteGuard<'_, BufferedIO>> {
4374 if !self.closed.load() {
4375 Ok(self.buffer.write())
4376 } else {
4377 Err(io_closed_error(vm))
4378 }
4379 }
4380 }
4381
4382 #[pyclass(flags(BASETYPE, HAS_DICT, HAS_WEAKREF), with(Constructor, Initializer))]
4383 impl StringIO {
4384 #[pymethod]
4385 const fn readable(&self) -> bool {
4386 true
4387 }
4388
4389 #[pymethod]
4390 const fn writable(&self) -> bool {
4391 true
4392 }
4393
4394 #[pymethod]
4395 const fn seekable(&self) -> bool {
4396 true
4397 }
4398
4399 #[pygetset]
4400 fn closed(&self) -> bool {
4401 self.closed.load()
4402 }
4403
4404 #[pymethod]
4405 fn close(&self) {
4406 self.closed.store(true);
4407 }
4408
4409 #[pymethod]
4411 fn write(&self, data: PyStrRef, vm: &VirtualMachine) -> PyResult<u64> {
4412 let bytes = data.as_bytes();
4413 self.buffer(vm)?
4414 .write(bytes)
4415 .ok_or_else(|| vm.new_type_error("Error Writing String"))
4416 }
4417
4418 #[pymethod]
4420 fn getvalue(&self, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
4421 let bytes = self.buffer(vm)?.getvalue();
4422 Wtf8Buf::from_bytes(bytes).map_err(|_| vm.new_value_error("Error Retrieving Value"))
4423 }
4424
4425 #[pymethod]
4427 fn seek(
4428 &self,
4429 offset: PyObjectRef,
4430 how: OptionalArg<i32>,
4431 vm: &VirtualMachine,
4432 ) -> PyResult<u64> {
4433 self.buffer(vm)?
4434 .seek(seekfrom(vm, offset, how)?)
4435 .map_err(|err| os_err(vm, err))
4436 }
4437
4438 #[pymethod]
4442 fn read(&self, size: OptionalSize, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
4443 let data = self.buffer(vm)?.read(size.to_usize()).unwrap_or_default();
4444
4445 let value = Wtf8Buf::from_bytes(data)
4446 .map_err(|_| vm.new_value_error("Error Retrieving Value"))?;
4447 Ok(value)
4448 }
4449
4450 #[pymethod]
4451 fn tell(&self, vm: &VirtualMachine) -> PyResult<u64> {
4452 Ok(self.buffer(vm)?.tell())
4453 }
4454
4455 #[pymethod]
4456 fn readline(&self, size: OptionalSize, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
4457 let input = self.buffer(vm)?.readline(size.to_usize(), vm)?;
4460 Wtf8Buf::from_bytes(input).map_err(|_| vm.new_value_error("Error Retrieving Value"))
4461 }
4462
4463 #[pymethod]
4464 fn truncate(&self, pos: OptionalSize, vm: &VirtualMachine) -> PyResult<usize> {
4465 let mut buffer = self.buffer(vm)?;
4466 let pos = pos.try_usize(vm)?;
4467 Ok(buffer.truncate(pos))
4468 }
4469
4470 #[pygetset]
4471 const fn line_buffering(&self) -> bool {
4472 false
4473 }
4474
4475 #[pymethod]
4476 fn __getstate__(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
4477 let buffer = zelf.buffer(vm)?;
4478 let content = Wtf8Buf::from_bytes(buffer.getvalue())
4479 .map_err(|_| vm.new_value_error("Error Retrieving Value"))?;
4480 let pos = buffer.tell();
4481 drop(buffer);
4482
4483 let dict_obj: PyObjectRef = match zelf.as_object().dict() {
4485 Some(d) if !d.is_empty() => d.into(),
4486 _ => vm.ctx.none(),
4487 };
4488
4489 Ok(vm.ctx.new_tuple(vec![
4492 vm.ctx.new_str(content).into(),
4493 vm.ctx.new_str("\n").into(),
4494 vm.ctx.new_int(pos).into(),
4495 dict_obj,
4496 ]))
4497 }
4498
4499 #[pymethod]
4500 fn __setstate__(zelf: PyRef<Self>, state: PyTupleRef, vm: &VirtualMachine) -> PyResult<()> {
4501 if zelf.closed.load() {
4503 return Err(vm.new_value_error("__setstate__ on closed file"));
4504 }
4505 if state.len() != 4 {
4506 return Err(vm.new_type_error(format!(
4507 "__setstate__ argument should be 4-tuple, got {}",
4508 state.len()
4509 )));
4510 }
4511
4512 let content: PyStrRef = state[0].clone().try_into_value(vm)?;
4513 let pos: u64 = state[2].clone().try_into_value(vm)?;
4515 let dict = &state[3];
4516
4517 let raw_bytes = content.as_bytes().to_vec();
4519 let mut buffer = zelf.buffer.write();
4520 *buffer = BufferedIO::new(Cursor::new(raw_bytes));
4521 buffer
4522 .seek(SeekFrom::Start(pos))
4523 .map_err(|err| os_err(vm, err))?;
4524 drop(buffer);
4525
4526 if !vm.is_none(dict) {
4528 let dict_ref: PyRef<PyDict> = dict.clone().try_into_value(vm)?;
4529 if let Some(obj_dict) = zelf.as_object().dict() {
4530 obj_dict.clear();
4531 for (key, value) in dict_ref.into_iter() {
4532 obj_dict.set_item(&*key, value, vm)?;
4533 }
4534 }
4535 }
4536
4537 Ok(())
4538 }
4539 }
4540
4541 #[derive(FromArgs)]
4542 struct BytesIOArgs {
4543 #[pyarg(any, optional)]
4544 initial_bytes: OptionalArg<Option<ArgBytesLike>>,
4545 }
4546
4547 #[pyattr]
4548 #[pyclass(name = "BytesIO", base = _BufferedIOBase)]
4549 #[derive(Debug)]
4550 struct BytesIO {
4551 _base: _BufferedIOBase,
4552 buffer: PyRwLock<BufferedIO>,
4553 closed: AtomicCell<bool>,
4554 exports: AtomicCell<usize>,
4555 }
4556
4557 impl Constructor for BytesIO {
4558 type Args = FuncArgs;
4559
4560 fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
4561 Ok(Self {
4562 _base: Default::default(),
4563 buffer: PyRwLock::new(BufferedIO::new(Cursor::new(Vec::new()))),
4564 closed: AtomicCell::new(false),
4565 exports: AtomicCell::new(0),
4566 })
4567 }
4568 }
4569
4570 impl Initializer for BytesIO {
4571 type Args = BytesIOArgs;
4572
4573 fn init(zelf: PyRef<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> {
4574 if zelf.exports.load() > 0 {
4575 return Err(
4576 vm.new_buffer_error("Existing exports of data: object cannot be re-sized")
4577 );
4578 }
4579
4580 let raw_bytes = args
4581 .initial_bytes
4582 .flatten()
4583 .map_or_else(Vec::new, |input| input.borrow_buf().to_vec());
4584 *zelf.buffer.write() = BufferedIO::new(Cursor::new(raw_bytes));
4585 Ok(())
4586 }
4587 }
4588
4589 impl BytesIO {
4590 fn buffer(&self, vm: &VirtualMachine) -> PyResult<PyRwLockWriteGuard<'_, BufferedIO>> {
4591 if !self.closed.load() {
4592 Ok(self.buffer.write())
4593 } else {
4594 Err(io_closed_error(vm))
4595 }
4596 }
4597 }
4598
4599 #[pyclass(
4600 flags(BASETYPE, HAS_DICT, HAS_WEAKREF),
4601 with(PyRef, Constructor, Initializer)
4602 )]
4603 impl BytesIO {
4604 #[pymethod]
4605 const fn readable(&self) -> bool {
4606 true
4607 }
4608
4609 #[pymethod]
4610 const fn writable(&self) -> bool {
4611 true
4612 }
4613
4614 #[pymethod]
4615 const fn seekable(&self) -> bool {
4616 true
4617 }
4618
4619 #[pymethod]
4620 fn flush(&self, vm: &VirtualMachine) -> PyResult<()> {
4621 if self.closed.load() {
4622 Err(io_closed_error(vm))
4623 } else {
4624 Ok(())
4625 }
4626 }
4627
4628 #[pymethod]
4629 fn write(&self, data: ArgBytesLike, vm: &VirtualMachine) -> PyResult<u64> {
4630 let mut buffer = self.try_resizable(vm)?;
4631 data.with_ref(|b| buffer.write(b))
4632 .ok_or_else(|| vm.new_type_error("Error Writing Bytes"))
4633 }
4634
4635 #[pymethod]
4637 fn getvalue(&self, vm: &VirtualMachine) -> PyResult<PyBytesRef> {
4638 let bytes = self.buffer(vm)?.getvalue();
4639 Ok(vm.ctx.new_bytes(bytes))
4640 }
4641
4642 #[pymethod]
4646 #[pymethod(name = "read1")]
4647 fn read(&self, size: OptionalSize, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
4648 let buf = self.buffer(vm)?.read(size.to_usize()).unwrap_or_default();
4649 Ok(buf)
4650 }
4651
4652 #[pymethod]
4653 fn readinto(&self, obj: ArgMemoryBuffer, vm: &VirtualMachine) -> PyResult<usize> {
4654 let mut buf = self.buffer(vm)?;
4655 let ret = buf
4656 .cursor
4657 .read(&mut obj.borrow_buf_mut())
4658 .map_err(|_| vm.new_value_error("Error readinto from Take"))?;
4659
4660 Ok(ret)
4661 }
4662
4663 #[pymethod]
4665 fn seek(
4666 &self,
4667 offset: PyObjectRef,
4668 how: OptionalArg<i32>,
4669 vm: &VirtualMachine,
4670 ) -> PyResult<u64> {
4671 let seek_from = seekfrom(vm, offset, how)?;
4672 let mut buffer = self.buffer(vm)?;
4673
4674 match seek_from {
4676 SeekFrom::Current(offset) if offset < 0 => {
4677 let current = buffer.tell();
4678 let new_pos = current.saturating_add_signed(offset);
4679 buffer
4680 .seek(SeekFrom::Start(new_pos))
4681 .map_err(|err| os_err(vm, err))
4682 }
4683 _ => buffer.seek(seek_from).map_err(|err| os_err(vm, err)),
4684 }
4685 }
4686
4687 #[pymethod]
4688 fn tell(&self, vm: &VirtualMachine) -> PyResult<u64> {
4689 Ok(self.buffer(vm)?.tell())
4690 }
4691
4692 #[pymethod]
4693 fn readline(&self, size: OptionalSize, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
4694 self.buffer(vm)?.readline(size.to_usize(), vm)
4695 }
4696
4697 #[pymethod]
4698 fn truncate(&self, pos: OptionalSize, vm: &VirtualMachine) -> PyResult<usize> {
4699 if self.closed.load() {
4700 return Err(io_closed_error(vm));
4701 }
4702 let mut buffer = self.try_resizable(vm)?;
4703 let pos = pos.try_usize(vm)?;
4704 Ok(buffer.truncate(pos))
4705 }
4706
4707 #[pygetset]
4708 fn closed(&self) -> bool {
4709 self.closed.load()
4710 }
4711
4712 #[pymethod]
4713 fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
4714 if self.exports.load() > 0 {
4715 return Err(
4716 vm.new_buffer_error("Existing exports of data: object cannot be closed")
4717 );
4718 }
4719 self.closed.store(true);
4720 Ok(())
4721 }
4722
4723 #[pymethod]
4724 fn __getstate__(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
4725 let buffer = zelf.buffer(vm)?;
4726 let content = buffer.getvalue();
4727 let pos = buffer.tell();
4728 drop(buffer);
4729
4730 let dict_obj: PyObjectRef = match zelf.as_object().dict() {
4732 Some(d) if !d.is_empty() => d.into(),
4733 _ => vm.ctx.none(),
4734 };
4735
4736 Ok(vm.ctx.new_tuple(vec![
4738 vm.ctx.new_bytes(content).into(),
4739 vm.ctx.new_int(pos).into(),
4740 dict_obj,
4741 ]))
4742 }
4743
4744 #[pymethod]
4745 fn __setstate__(zelf: PyRef<Self>, state: PyTupleRef, vm: &VirtualMachine) -> PyResult<()> {
4746 if zelf.closed.load() {
4747 return Err(vm.new_value_error("__setstate__ on closed file"));
4748 }
4749 if state.len() != 3 {
4750 return Err(vm.new_type_error(format!(
4751 "__setstate__ argument should be 3-tuple, got {}",
4752 state.len()
4753 )));
4754 }
4755
4756 let content: PyBytesRef = state[0].clone().try_into_value(vm)?;
4757 let pos: u64 = state[1].clone().try_into_value(vm)?;
4758 let dict = &state[2];
4759
4760 let mut buffer = zelf.try_resizable(vm)?;
4762 *buffer = BufferedIO::new(Cursor::new(content.as_bytes().to_vec()));
4763 buffer
4764 .seek(SeekFrom::Start(pos))
4765 .map_err(|err| os_err(vm, err))?;
4766 drop(buffer);
4767
4768 if !vm.is_none(dict) {
4770 let dict_ref: PyRef<PyDict> = dict.clone().try_into_value(vm)?;
4771 if let Some(obj_dict) = zelf.as_object().dict() {
4772 obj_dict.clear();
4773 for (key, value) in dict_ref.into_iter() {
4774 obj_dict.set_item(&*key, value, vm)?;
4775 }
4776 }
4777 }
4778
4779 Ok(())
4780 }
4781
4782 #[pymethod]
4783 fn isatty(&self, vm: &VirtualMachine) -> PyResult<bool> {
4784 if self.closed() {
4785 return Err(io_closed_error(vm));
4786 }
4787
4788 Ok(false)
4789 }
4790 }
4791
4792 #[pyclass]
4793 impl PyRef<BytesIO> {
4794 #[pymethod]
4795 fn getbuffer(self, vm: &VirtualMachine) -> PyResult<PyMemoryView> {
4796 if self.closed.load() {
4797 return Err(vm.new_value_error("I/O operation on closed file."));
4798 }
4799 let len = self.buffer.read().cursor.get_ref().len();
4800 let buffer = PyBuffer::new(
4801 self.into(),
4802 BufferDescriptor::simple(len, false),
4803 &BYTES_IO_BUFFER_METHODS,
4804 );
4805 let view = PyMemoryView::from_buffer(buffer, vm)?;
4806 Ok(view)
4807 }
4808 }
4809
4810 static BYTES_IO_BUFFER_METHODS: BufferMethods = BufferMethods {
4811 obj_bytes: |buffer| {
4812 let zelf = buffer.obj_as::<BytesIO>();
4813 PyRwLockReadGuard::map(zelf.buffer.read(), |x| x.cursor.get_ref().as_slice()).into()
4814 },
4815 obj_bytes_mut: |buffer| {
4816 let zelf = buffer.obj_as::<BytesIO>();
4817 PyRwLockWriteGuard::map(zelf.buffer.write(), |x| x.cursor.get_mut().as_mut_slice())
4818 .into()
4819 },
4820
4821 release: |buffer| {
4822 buffer.obj_as::<BytesIO>().exports.fetch_sub(1);
4823 },
4824
4825 retain: |buffer| {
4826 buffer.obj_as::<BytesIO>().exports.fetch_add(1);
4827 },
4828 };
4829
4830 impl BufferResizeGuard for BytesIO {
4831 type Resizable<'a> = PyRwLockWriteGuard<'a, BufferedIO>;
4832
4833 fn try_resizable_opt(&self) -> Option<Self::Resizable<'_>> {
4834 let w = self.buffer.write();
4835 (self.exports.load() == 0).then_some(w)
4836 }
4837 }
4838
4839 #[repr(u8)]
4840 #[derive(Debug)]
4841 enum FileMode {
4842 Read = b'r',
4843 Write = b'w',
4844 Exclusive = b'x',
4845 Append = b'a',
4846 }
4847
4848 #[repr(u8)]
4849 #[derive(Debug)]
4850 enum EncodeMode {
4851 Text = b't',
4852 Bytes = b'b',
4853 }
4854
4855 #[derive(Debug)]
4856 struct Mode {
4857 file: FileMode,
4858 encode: EncodeMode,
4859 plus: bool,
4860 }
4861
4862 impl core::str::FromStr for Mode {
4863 type Err = ParseModeError;
4864
4865 fn from_str(s: &str) -> Result<Self, Self::Err> {
4866 let mut file = None;
4867 let mut encode = None;
4868 let mut plus = false;
4869 macro_rules! set_mode {
4870 ($var:ident, $mode:path, $err:ident) => {{
4871 match $var {
4872 Some($mode) => return Err(ParseModeError::InvalidMode),
4873 Some(_) => return Err(ParseModeError::$err),
4874 None => $var = Some($mode),
4875 }
4876 }};
4877 }
4878
4879 for ch in s.chars() {
4880 match ch {
4881 '+' => {
4882 if plus {
4883 return Err(ParseModeError::InvalidMode);
4884 }
4885 plus = true
4886 }
4887 't' => set_mode!(encode, EncodeMode::Text, MultipleEncode),
4888 'b' => set_mode!(encode, EncodeMode::Bytes, MultipleEncode),
4889 'r' => set_mode!(file, FileMode::Read, MultipleFile),
4890 'a' => set_mode!(file, FileMode::Append, MultipleFile),
4891 'w' => set_mode!(file, FileMode::Write, MultipleFile),
4892 'x' => set_mode!(file, FileMode::Exclusive, MultipleFile),
4893 _ => return Err(ParseModeError::InvalidMode),
4894 }
4895 }
4896
4897 let file = file.ok_or(ParseModeError::NoFile)?;
4898 let encode = encode.unwrap_or(EncodeMode::Text);
4899
4900 Ok(Self { file, encode, plus })
4901 }
4902 }
4903
4904 impl Mode {
4905 const fn rawmode(&self) -> &'static str {
4906 match (&self.file, self.plus) {
4907 (FileMode::Read, true) => "rb+",
4908 (FileMode::Read, false) => "rb",
4909 (FileMode::Write, true) => "wb+",
4910 (FileMode::Write, false) => "wb",
4911 (FileMode::Exclusive, true) => "xb+",
4912 (FileMode::Exclusive, false) => "xb",
4913 (FileMode::Append, true) => "ab+",
4914 (FileMode::Append, false) => "ab",
4915 }
4916 }
4917 }
4918
4919 enum ParseModeError {
4920 InvalidMode,
4921 MultipleFile,
4922 MultipleEncode,
4923 NoFile,
4924 }
4925
4926 impl ParseModeError {
4927 fn error_msg(&self, mode_string: &str) -> String {
4928 match self {
4929 Self::InvalidMode => format!("invalid mode: '{mode_string}'"),
4930 Self::MultipleFile => {
4931 "must have exactly one of create/read/write/append mode".to_owned()
4932 }
4933 Self::MultipleEncode => "can't have text and binary mode at once".to_owned(),
4934 Self::NoFile => {
4935 "Must have exactly one of create/read/write/append mode and at most one plus"
4936 .to_owned()
4937 }
4938 }
4939 }
4940 }
4941
4942 #[derive(FromArgs)]
4943 struct IoOpenArgs {
4944 file: PyObjectRef,
4945 #[pyarg(any, optional)]
4946 mode: OptionalArg<PyUtf8StrRef>,
4947 #[pyarg(flatten)]
4948 opts: OpenArgs,
4949 }
4950
4951 #[pyfunction]
4952 fn open(args: IoOpenArgs, vm: &VirtualMachine) -> PyResult {
4953 io_open(
4954 args.file,
4955 args.mode.as_ref().into_option().map(|s| s.as_str()),
4956 args.opts,
4957 vm,
4958 )
4959 }
4960
4961 #[pyfunction]
4962 fn open_code(file: PyObjectRef, vm: &VirtualMachine) -> PyResult {
4963 io_open(file, Some("rb"), OpenArgs::default(), vm)
4965 }
4966
4967 #[derive(FromArgs)]
4968 pub struct OpenArgs {
4969 #[pyarg(any, default = -1)]
4970 pub buffering: isize,
4971 #[pyarg(any, default)]
4972 pub encoding: Option<PyUtf8StrRef>,
4973 #[pyarg(any, default)]
4974 pub errors: Option<PyUtf8StrRef>,
4975 #[pyarg(any, default)]
4976 pub newline: Option<PyUtf8StrRef>,
4977 #[pyarg(any, default = true)]
4978 pub closefd: bool,
4979 #[pyarg(any, default)]
4980 pub opener: Option<PyObjectRef>,
4981 }
4982
4983 impl Default for OpenArgs {
4984 fn default() -> Self {
4985 Self {
4986 buffering: -1,
4987 encoding: None,
4988 errors: None,
4989 newline: None,
4990 closefd: true,
4991 opener: None,
4992 }
4993 }
4994 }
4995
4996 #[cfg(all(unix, feature = "threading"))]
5003 pub unsafe fn reinit_std_streams_after_fork(vm: &VirtualMachine) {
5004 for name in ["stdin", "stdout", "stderr"] {
5005 let Ok(stream) = vm.sys_module.get_attr(name, vm) else {
5006 continue;
5007 };
5008 reinit_io_locks(&stream);
5009 }
5010 }
5011
5012 #[cfg(all(unix, feature = "threading"))]
5013 fn reinit_io_locks(obj: &PyObject) {
5014 use crate::common::lock::reinit_thread_mutex_after_fork;
5015
5016 if let Some(tio) = obj.downcast_ref::<TextIOWrapper>() {
5017 unsafe { reinit_thread_mutex_after_fork(&tio.data) };
5018 if let Some(guard) = tio.data.lock()
5019 && let Some(ref data) = *guard
5020 {
5021 if let Some(ref decoder) = data.decoder {
5022 reinit_io_locks(decoder);
5023 }
5024 reinit_io_locks(&data.buffer);
5025 }
5026 return;
5027 }
5028 if let Some(nl) = obj.downcast_ref::<IncrementalNewlineDecoder>() {
5029 unsafe { reinit_thread_mutex_after_fork(&nl.data) };
5030 return;
5031 }
5032 if let Some(br) = obj.downcast_ref::<BufferedReader>() {
5033 unsafe { reinit_thread_mutex_after_fork(&br.data) };
5034 return;
5035 }
5036 if let Some(bw) = obj.downcast_ref::<BufferedWriter>() {
5037 unsafe { reinit_thread_mutex_after_fork(&bw.data) };
5038 return;
5039 }
5040 if let Some(brw) = obj.downcast_ref::<BufferedRandom>() {
5041 unsafe { reinit_thread_mutex_after_fork(&brw.data) };
5042 return;
5043 }
5044 if let Some(brw) = obj.downcast_ref::<BufferedRWPair>() {
5045 unsafe { reinit_thread_mutex_after_fork(&brw.read.data) };
5046 unsafe { reinit_thread_mutex_after_fork(&brw.write.data) };
5047 }
5048 }
5049
5050 pub fn io_open(
5051 file: PyObjectRef,
5052 mode: Option<&str>,
5053 opts: OpenArgs,
5054 vm: &VirtualMachine,
5055 ) -> PyResult {
5056 let mode_string = mode.unwrap_or("r");
5058 let mode = mode_string
5059 .parse::<Mode>()
5060 .map_err(|e| vm.new_value_error(e.error_msg(mode_string)))?;
5061
5062 if let EncodeMode::Bytes = mode.encode {
5063 let msg = if opts.encoding.is_some() {
5064 Some("binary mode doesn't take an encoding argument")
5065 } else if opts.errors.is_some() {
5066 Some("binary mode doesn't take an errors argument")
5067 } else if opts.newline.is_some() {
5068 Some("binary mode doesn't take a newline argument")
5069 } else {
5070 None
5071 };
5072 if let Some(msg) = msg {
5073 return Err(vm.new_value_error(msg));
5074 }
5075 }
5076
5077 #[cfg(all(unix, feature = "host_env"))]
5079 if let Ok(crate::ospath::OsPathOrFd::Fd(fd)) = file.clone().try_into_value(vm) {
5080 nix::fcntl::fcntl(fd, nix::fcntl::F_GETFD).map_err(|_| vm.new_last_errno_error())?;
5081 }
5082
5083 #[cfg(all(feature = "host_env", windows))]
5087 let is_console = super::winconsoleio::pyio_get_console_type(&file, vm) != '\0';
5088 #[cfg(not(all(feature = "host_env", windows)))]
5089 let is_console = false;
5090
5091 let file_io_class: &Py<PyType> = {
5092 cfg_if::cfg_if! {
5093 if #[cfg(all(feature = "host_env", windows))] {
5094 if is_console {
5095 Some(super::winconsoleio::WindowsConsoleIO::static_type())
5096 } else {
5097 Some(super::fileio::FileIO::static_type())
5098 }
5099 } else if #[cfg(feature = "host_env")] {
5100 Some(super::fileio::FileIO::static_type())
5101 } else {
5102 None
5103 }
5104 }
5105 }
5106 .ok_or_else(|| {
5107 new_unsupported_operation(
5108 vm,
5109 "Couldn't get FileIO, io.open likely isn't supported on your platform".to_owned(),
5110 )
5111 })?;
5112 let raw = PyType::call(
5113 file_io_class,
5114 (file, mode.rawmode(), opts.closefd, opts.opener).into_args(vm),
5115 vm,
5116 )?;
5117
5118 let isatty = opts.buffering < 0 && {
5119 let atty = vm.call_method(&raw, "isatty", ())?;
5120 bool::try_from_object(vm, atty)?
5121 };
5122
5123 if opts.buffering == 1 && matches!(mode.encode, EncodeMode::Bytes) {
5125 crate::stdlib::_warnings::warn(
5126 vm.ctx.exceptions.runtime_warning,
5127 "line buffering (buffering=1) isn't supported in binary mode, the default buffer size will be used".to_owned(),
5128 1,
5129 vm,
5130 )?;
5131 }
5132
5133 let line_buffering = opts.buffering == 1 || isatty;
5134
5135 let buffering = if opts.buffering < 0 || opts.buffering == 1 {
5136 DEFAULT_BUFFER_SIZE
5137 } else {
5138 opts.buffering as usize
5139 };
5140
5141 if buffering == 0 {
5142 let ret = match mode.encode {
5143 EncodeMode::Text => {
5144 let _ = vm.call_method(&raw, "close", ());
5145 Err(vm.new_value_error("can't have unbuffered text I/O"))
5146 }
5147 EncodeMode::Bytes => Ok(raw),
5148 };
5149 return ret;
5150 }
5151
5152 let cls = if mode.plus {
5153 BufferedRandom::static_type()
5154 } else if let FileMode::Read = mode.file {
5155 BufferedReader::static_type()
5156 } else {
5157 BufferedWriter::static_type()
5158 };
5159 let buffered = PyType::call(cls, (raw, buffering).into_args(vm), vm)?;
5160
5161 match mode.encode {
5162 EncodeMode::Text => {
5163 let encoding = if is_console && opts.encoding.is_none() {
5164 Some(PyUtf8Str::from("utf-8").into_ref(&vm.ctx))
5166 } else {
5167 match opts.encoding {
5168 Some(enc) => Some(enc),
5169 None => {
5170 let encoding =
5171 text_encoding(vm.ctx.none(), OptionalArg::Present(2), vm)?;
5172 Some(PyUtf8StrRef::try_from_object(vm, encoding.into())?)
5173 }
5174 }
5175 };
5176 let tio = TextIOWrapper::static_type();
5177 let wrapper = PyType::call(
5178 tio,
5179 (
5180 buffered.clone(),
5181 encoding,
5182 opts.errors,
5183 opts.newline,
5184 line_buffering,
5185 )
5186 .into_args(vm),
5187 vm,
5188 )
5189 .inspect_err(|_err| {
5190 let _ = vm.call_method(&buffered, "close", ());
5191 })?;
5192 wrapper.set_attr("mode", vm.new_pyobj(mode_string), vm)?;
5193 Ok(wrapper)
5194 }
5195 EncodeMode::Bytes => Ok(buffered),
5196 }
5197 }
5198
5199 fn create_unsupported_operation(ctx: &Context) -> PyTypeRef {
5200 use crate::builtins::type_::PyAttributes;
5201 use crate::types::PyTypeSlots;
5202
5203 let mut attrs = PyAttributes::default();
5204 attrs.insert(identifier!(ctx, __module__), ctx.new_str("io").into());
5205
5206 PyType::new_heap(
5207 "UnsupportedOperation",
5208 vec![
5209 ctx.exceptions.os_error.to_owned(),
5210 ctx.exceptions.value_error.to_owned(),
5211 ],
5212 attrs,
5213 PyTypeSlots::heap_default(),
5214 ctx.types.type_type.to_owned(),
5215 ctx,
5216 )
5217 .unwrap()
5218 }
5219
5220 pub fn unsupported_operation() -> &'static Py<PyType> {
5221 rustpython_common::static_cell! {
5222 static CELL: PyTypeRef;
5223 }
5224 CELL.get_or_init(|| create_unsupported_operation(Context::genesis()))
5225 }
5226
5227 #[pyfunction]
5228 fn text_encoding(
5229 encoding: PyObjectRef,
5230 stacklevel: OptionalArg<i32>,
5231 vm: &VirtualMachine,
5232 ) -> PyResult<PyStrRef> {
5233 if vm.is_none(&encoding) {
5234 let encoding = if vm.state.config.settings.utf8_mode > 0 {
5235 "utf-8"
5236 } else {
5237 "locale"
5238 };
5239 if vm.state.config.settings.warn_default_encoding {
5240 let mut stacklevel = stacklevel.unwrap_or(2);
5241 if stacklevel > 1
5242 && let Some(frame) = vm.current_frame()
5243 && let Some(stdlib_dir) = vm.state.config.paths.stdlib_dir.as_deref()
5244 {
5245 let path = frame.code.source_path().as_str();
5246 if !path.starts_with(stdlib_dir) {
5247 stacklevel = stacklevel.saturating_sub(1);
5248 }
5249 }
5250 let stacklevel = usize::try_from(stacklevel).unwrap_or(0);
5251 crate::stdlib::_warnings::warn(
5252 vm.ctx.exceptions.encoding_warning,
5253 "'encoding' argument not specified".to_owned(),
5254 stacklevel,
5255 vm,
5256 )?;
5257 }
5258 return Ok(vm.ctx.new_str(encoding));
5259 }
5260 encoding.try_into_value(vm)
5261 }
5262
5263 #[cfg(test)]
5264 mod tests {
5265 use super::*;
5266
5267 #[test]
5268 fn test_buffered_read() {
5269 let data = vec![1, 2, 3, 4];
5270 let bytes = None;
5271 let mut buffered = BufferedIO {
5272 cursor: Cursor::new(data.clone()),
5273 };
5274
5275 assert_eq!(buffered.read(bytes).unwrap(), data);
5276 }
5277
5278 #[test]
5279 fn test_buffered_seek() {
5280 let data = vec![1, 2, 3, 4];
5281 let count: u64 = 2;
5282 let mut buffered = BufferedIO {
5283 cursor: Cursor::new(data),
5284 };
5285
5286 assert_eq!(buffered.seek(SeekFrom::Start(count)).unwrap(), count);
5287 assert_eq!(buffered.read(Some(count as usize)).unwrap(), vec![3, 4]);
5288 }
5289
5290 #[test]
5291 fn test_buffered_value() {
5292 let data = vec![1, 2, 3, 4];
5293 let buffered = BufferedIO {
5294 cursor: Cursor::new(data.clone()),
5295 };
5296
5297 assert_eq!(buffered.getvalue(), data);
5298 }
5299 }
5300
5301 pub(crate) fn module_exec(vm: &VirtualMachine, module: &Py<PyModule>) -> PyResult<()> {
5302 __module_exec(vm, module);
5304
5305 #[cfg(feature = "host_env")]
5307 super::fileio::module_exec(vm, module)?;
5308
5309 #[cfg(all(feature = "host_env", windows))]
5311 super::winconsoleio::module_exec(vm, module)?;
5312
5313 let unsupported_operation = unsupported_operation().to_owned();
5314 extend_module!(vm, module, {
5315 "UnsupportedOperation" => unsupported_operation,
5316 "BlockingIOError" => vm.ctx.exceptions.blocking_io_error.to_owned(),
5317 });
5318 Ok(())
5319 }
5320}
5321#[cfg(feature = "host_env")]
5323#[pymodule]
5324mod fileio {
5325 use super::{_io::*, Offset, iobase_finalize};
5326 use crate::{
5327 AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
5328 VirtualMachine,
5329 builtins::{PyBaseExceptionRef, PyUtf8Str, PyUtf8StrRef},
5330 common::{crt_fd, wtf8::Wtf8Buf},
5331 convert::{IntoPyException, ToPyException},
5332 exceptions::OSErrorBuilder,
5333 function::{ArgBytesLike, ArgMemoryBuffer, OptionalArg, OptionalOption},
5334 ospath::{OsPath, OsPathOrFd},
5335 stdlib::os,
5336 types::{Constructor, DefaultConstructor, Destructor, Initializer, Representable},
5337 };
5338 use crossbeam_utils::atomic::AtomicCell;
5339 use std::io::Read;
5340
5341 bitflags::bitflags! {
5342 #[derive(Copy, Clone, Debug, PartialEq)]
5343 struct Mode: u8 {
5344 const CREATED = 0b0001;
5345 const READABLE = 0b0010;
5346 const WRITABLE = 0b0100;
5347 const APPENDING = 0b1000;
5348 }
5349 }
5350
5351 enum ModeError {
5352 Invalid,
5353 BadRwa,
5354 }
5355
5356 impl ModeError {
5357 fn error_msg(&self, mode_str: &str) -> String {
5358 match self {
5359 Self::Invalid => format!("invalid mode: {mode_str}"),
5360 Self::BadRwa => {
5361 "Must have exactly one of create/read/write/append mode and at most one plus"
5362 .to_owned()
5363 }
5364 }
5365 }
5366 }
5367
5368 fn compute_mode(mode_str: &str) -> Result<(Mode, i32), ModeError> {
5369 let mut flags = 0;
5370 let mut plus = false;
5371 let mut rwa = false;
5372 let mut mode = Mode::empty();
5373 for c in mode_str.bytes() {
5374 match c {
5375 b'x' => {
5376 if rwa {
5377 return Err(ModeError::BadRwa);
5378 }
5379 rwa = true;
5380 mode.insert(Mode::WRITABLE | Mode::CREATED);
5381 flags |= libc::O_EXCL | libc::O_CREAT;
5382 }
5383 b'r' => {
5384 if rwa {
5385 return Err(ModeError::BadRwa);
5386 }
5387 rwa = true;
5388 mode.insert(Mode::READABLE);
5389 }
5390 b'w' => {
5391 if rwa {
5392 return Err(ModeError::BadRwa);
5393 }
5394 rwa = true;
5395 mode.insert(Mode::WRITABLE);
5396 flags |= libc::O_CREAT | libc::O_TRUNC;
5397 }
5398 b'a' => {
5399 if rwa {
5400 return Err(ModeError::BadRwa);
5401 }
5402 rwa = true;
5403 mode.insert(Mode::WRITABLE | Mode::APPENDING);
5404 flags |= libc::O_APPEND | libc::O_CREAT;
5405 }
5406 b'+' => {
5407 if plus {
5408 return Err(ModeError::BadRwa);
5409 }
5410 plus = true;
5411 mode.insert(Mode::READABLE | Mode::WRITABLE);
5412 }
5413 b'b' => {}
5414 _ => return Err(ModeError::Invalid),
5415 }
5416 }
5417
5418 if !rwa {
5419 return Err(ModeError::BadRwa);
5420 }
5421
5422 if mode.contains(Mode::READABLE | Mode::WRITABLE) {
5423 flags |= libc::O_RDWR
5424 } else if mode.contains(Mode::READABLE) {
5425 flags |= libc::O_RDONLY
5426 } else {
5427 flags |= libc::O_WRONLY
5428 }
5429
5430 #[cfg(windows)]
5431 {
5432 flags |= libc::O_BINARY | libc::O_NOINHERIT;
5433 }
5434 #[cfg(unix)]
5435 {
5436 flags |= libc::O_CLOEXEC
5437 }
5438
5439 Ok((mode, flags as _))
5440 }
5441
5442 #[pyattr]
5443 #[pyclass(module = "_io", name, base = _RawIOBase)]
5444 #[derive(Debug)]
5445 pub(super) struct FileIO {
5446 _base: _RawIOBase,
5447 fd: AtomicCell<i32>,
5448 closefd: AtomicCell<bool>,
5449 mode: AtomicCell<Mode>,
5450 seekable: AtomicCell<Option<bool>>,
5451 blksize: AtomicCell<i64>,
5452 finalizing: AtomicCell<bool>,
5453 }
5454
5455 #[derive(FromArgs)]
5456 pub struct FileIOArgs {
5457 #[pyarg(positional)]
5458 name: PyObjectRef,
5459 #[pyarg(any, default)]
5460 mode: Option<PyUtf8StrRef>,
5461 #[pyarg(any, default = true)]
5462 closefd: bool,
5463 #[pyarg(any, default)]
5464 opener: Option<PyObjectRef>,
5465 }
5466
5467 impl Default for FileIO {
5468 fn default() -> Self {
5469 Self {
5470 _base: Default::default(),
5471 fd: AtomicCell::new(-1),
5472 closefd: AtomicCell::new(true),
5473 mode: AtomicCell::new(Mode::empty()),
5474 seekable: AtomicCell::new(None),
5475 blksize: AtomicCell::new(super::DEFAULT_BUFFER_SIZE as _),
5476 finalizing: AtomicCell::new(false),
5477 }
5478 }
5479 }
5480
5481 impl DefaultConstructor for FileIO {}
5482
5483 impl Initializer for FileIO {
5484 type Args = FileIOArgs;
5485
5486 fn init(zelf: PyRef<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> {
5487 let name = args.name;
5489 if name.class().is(vm.ctx.types.bool_type) {
5491 crate::stdlib::_warnings::warn(
5492 vm.ctx.exceptions.runtime_warning,
5493 "bool is used as a file descriptor".to_owned(),
5494 1,
5495 vm,
5496 )?;
5497 }
5498 let arg_fd = if let Some(i) = name.downcast_ref::<crate::builtins::PyInt>() {
5499 let fd = i.try_to_primitive(vm)?;
5500 if fd < 0 {
5501 return Err(vm.new_value_error("negative file descriptor"));
5502 }
5503 Some(fd)
5504 } else {
5505 None
5506 };
5507
5508 let mode_obj = args
5509 .mode
5510 .unwrap_or_else(|| PyUtf8Str::from("rb").into_ref(&vm.ctx));
5511 let mode_str = mode_obj.as_str();
5512 let (mode, flags) =
5513 compute_mode(mode_str).map_err(|e| vm.new_value_error(e.error_msg(mode_str)))?;
5514 zelf.mode.store(mode);
5515
5516 let (fd, filename) = if let Some(fd) = arg_fd {
5517 zelf.closefd.store(args.closefd);
5518 (fd, None)
5519 } else {
5520 zelf.closefd.store(true);
5521 if !args.closefd {
5522 return Err(vm.new_value_error("Cannot use closefd=False with file name"));
5523 }
5524
5525 if let Some(opener) = args.opener {
5526 let fd = opener.call((name.clone(), flags), vm)?;
5527 if !fd.fast_isinstance(vm.ctx.types.int_type) {
5528 return Err(vm.new_type_error("expected integer from opener"));
5529 }
5530 let fd = i32::try_from_object(vm, fd)?;
5531 if fd < 0 {
5532 return Err(vm.new_value_error(format!("opener returned {fd}")));
5533 }
5534 (fd, None)
5535 } else {
5536 let path = OsPath::try_from_fspath(name.clone(), vm)?;
5537 #[cfg(any(unix, target_os = "wasi"))]
5538 let fd = crt_fd::open(&path.clone().into_cstring(vm)?, flags, 0o666);
5539 #[cfg(windows)]
5540 let fd = crt_fd::wopen(&path.to_wide_cstring(vm)?, flags, 0o666);
5541 let filename = OsPathOrFd::Path(path);
5542 match fd {
5543 Ok(fd) => (fd.into_raw(), Some(filename)),
5544 Err(e) => {
5545 return Err(OSErrorBuilder::with_filename_from_errno(&e, filename, vm));
5546 }
5547 }
5548 }
5549 };
5550 let fd_is_own = arg_fd.is_none();
5551 zelf.fd.store(fd);
5552 let fd = unsafe { crt_fd::Borrowed::borrow_raw(fd) };
5553 let filename = filename.unwrap_or(OsPathOrFd::Fd(fd));
5554
5555 let fd_fstat = crate::common::fileutils::fstat(fd);
5558
5559 #[cfg(windows)]
5560 {
5561 if let Err(err) = fd_fstat {
5562 if err.raw_os_error()
5564 == Some(windows_sys::Win32::Foundation::ERROR_INVALID_HANDLE as i32)
5565 {
5566 zelf.fd.store(-1);
5567 }
5568 return Err(OSErrorBuilder::with_filename(&err, filename, vm));
5569 }
5570 }
5571 #[cfg(any(unix, target_os = "wasi"))]
5572 {
5573 match fd_fstat {
5574 Ok(status) => {
5575 if (status.st_mode & libc::S_IFMT) == libc::S_IFDIR {
5576 if !fd_is_own {
5578 zelf.fd.store(-1);
5579 }
5580 let err = std::io::Error::from_raw_os_error(libc::EISDIR);
5581 return Err(OSErrorBuilder::with_filename(&err, filename, vm));
5582 }
5583 if status.st_blksize > 1 {
5585 #[allow(
5586 clippy::useless_conversion,
5587 reason = "needed for 32-bit platforms"
5588 )]
5589 zelf.blksize.store(i64::from(status.st_blksize));
5590 }
5591 }
5592 Err(err) => {
5593 if err.raw_os_error() == Some(libc::EBADF) {
5594 zelf.fd.store(-1);
5596 return Err(OSErrorBuilder::with_filename(&err, filename, vm));
5597 }
5598 }
5599 }
5600 }
5601
5602 #[cfg(windows)]
5603 crate::stdlib::msvcrt::setmode_binary(fd);
5604 if let Err(e) = zelf.as_object().set_attr("name", name, vm) {
5605 if !fd_is_own {
5607 zelf.fd.store(-1);
5608 }
5609 return Err(e);
5610 }
5611
5612 if mode.contains(Mode::APPENDING) {
5613 let _ = os::lseek(fd, 0, libc::SEEK_END, vm);
5614 }
5615
5616 Ok(())
5617 }
5618 }
5619
5620 impl Representable for FileIO {
5621 #[inline]
5622 fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
5623 let type_name = zelf.class().slot_name();
5624 let fd = zelf.fd.load();
5625 if fd < 0 {
5626 return Ok(format!("<{type_name} [closed]>"));
5627 }
5628 let name_repr = repr_file_obj_name(zelf.as_object(), vm)?;
5629 let mode = zelf.mode();
5630 let closefd = if zelf.closefd.load() { "True" } else { "False" };
5631 let repr = if let Some(name_repr) = name_repr {
5632 format!("<{type_name} name={name_repr} mode='{mode}' closefd={closefd}>")
5633 } else {
5634 format!("<{type_name} fd={fd} mode='{mode}' closefd={closefd}>")
5635 };
5636 Ok(repr)
5637 }
5638 }
5639
5640 #[pyclass(
5641 with(Constructor, Initializer, Representable, Destructor),
5642 flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
5643 )]
5644 impl FileIO {
5645 fn io_error(
5646 zelf: &Py<Self>,
5647 error: std::io::Error,
5648 vm: &VirtualMachine,
5649 ) -> PyBaseExceptionRef {
5650 let exc = error.to_pyexception(vm);
5651 if let Ok(name) = zelf.as_object().get_attr("name", vm) {
5652 exc.as_object()
5653 .set_attr("filename", name, vm)
5654 .expect("OSError.filename set must success");
5655 }
5656 exc
5657 }
5658
5659 #[pygetset]
5660 fn closed(&self) -> bool {
5661 self.fd.load() < 0
5662 }
5663
5664 #[pygetset]
5665 fn closefd(&self) -> bool {
5666 self.closefd.load()
5667 }
5668
5669 #[pygetset(name = "_blksize")]
5670 fn blksize(&self) -> i64 {
5671 self.blksize.load()
5672 }
5673
5674 #[pymethod]
5675 fn fileno(&self, vm: &VirtualMachine) -> PyResult<i32> {
5676 let fd = self.fd.load();
5677 if fd >= 0 {
5678 Ok(fd)
5679 } else {
5680 Err(io_closed_error(vm))
5681 }
5682 }
5683
5684 fn get_fd(&self, vm: &VirtualMachine) -> PyResult<crt_fd::Borrowed<'_>> {
5685 self.fileno(vm)
5686 .map(|fd| unsafe { crt_fd::Borrowed::borrow_raw(fd) })
5687 }
5688
5689 #[pymethod]
5690 fn readable(&self, vm: &VirtualMachine) -> PyResult<bool> {
5691 if self.fd.load() < 0 {
5692 return Err(io_closed_error(vm));
5693 }
5694 Ok(self.mode.load().contains(Mode::READABLE))
5695 }
5696
5697 #[pymethod]
5698 fn writable(&self, vm: &VirtualMachine) -> PyResult<bool> {
5699 if self.fd.load() < 0 {
5700 return Err(io_closed_error(vm));
5701 }
5702 Ok(self.mode.load().contains(Mode::WRITABLE))
5703 }
5704
5705 #[pygetset]
5706 fn mode(&self) -> &'static str {
5707 let mode = self.mode.load();
5708 if mode.contains(Mode::CREATED) {
5709 if mode.contains(Mode::READABLE) {
5710 "xb+"
5711 } else {
5712 "xb"
5713 }
5714 } else if mode.contains(Mode::APPENDING) {
5715 if mode.contains(Mode::READABLE) {
5716 "ab+"
5717 } else {
5718 "ab"
5719 }
5720 } else if mode.contains(Mode::READABLE) {
5721 if mode.contains(Mode::WRITABLE) {
5722 "rb+"
5723 } else {
5724 "rb"
5725 }
5726 } else {
5727 "wb"
5728 }
5729 }
5730
5731 #[pymethod]
5732 fn read(
5733 zelf: &Py<Self>,
5734 read_byte: OptionalSize,
5735 vm: &VirtualMachine,
5736 ) -> PyResult<Option<Vec<u8>>> {
5737 if !zelf.mode.load().contains(Mode::READABLE) {
5738 return Err(new_unsupported_operation(
5739 vm,
5740 "File or stream is not readable".to_owned(),
5741 ));
5742 }
5743 let handle = zelf.get_fd(vm)?;
5744 let bytes = if let Some(read_byte) = read_byte.to_usize() {
5745 let mut bytes = vec![0; read_byte];
5746 let n = loop {
5748 match vm.allow_threads(|| crt_fd::read(handle, &mut bytes)) {
5749 Ok(n) => break n,
5750 Err(e) if e.raw_os_error() == Some(libc::EINTR) => {
5751 vm.check_signals()?;
5752 continue;
5753 }
5754 Err(e) if e.raw_os_error() == Some(libc::EAGAIN) => {
5756 return Ok(None);
5757 }
5758 Err(e) => return Err(Self::io_error(zelf, e, vm)),
5759 }
5760 };
5761 bytes.truncate(n);
5762 bytes
5763 } else {
5764 let mut bytes = vec![];
5765 loop {
5767 match vm.allow_threads(|| {
5768 let mut h = handle;
5769 h.read_to_end(&mut bytes)
5770 }) {
5771 Ok(_) => break,
5772 Err(e) if e.raw_os_error() == Some(libc::EINTR) => {
5773 vm.check_signals()?;
5774 continue;
5775 }
5776 Err(e) if e.raw_os_error() == Some(libc::EAGAIN) => {
5778 if bytes.is_empty() {
5779 return Ok(None);
5780 }
5781 break;
5782 }
5783 Err(e) => return Err(Self::io_error(zelf, e, vm)),
5784 }
5785 }
5786 bytes
5787 };
5788
5789 Ok(Some(bytes))
5790 }
5791
5792 #[pymethod]
5793 fn readinto(
5794 zelf: &Py<Self>,
5795 obj: ArgMemoryBuffer,
5796 vm: &VirtualMachine,
5797 ) -> PyResult<Option<usize>> {
5798 if !zelf.mode.load().contains(Mode::READABLE) {
5799 return Err(new_unsupported_operation(
5800 vm,
5801 "File or stream is not readable".to_owned(),
5802 ));
5803 }
5804
5805 let handle = zelf.get_fd(vm)?;
5806
5807 let mut buf = obj.borrow_buf_mut();
5808 let ret = loop {
5810 match vm.allow_threads(|| crt_fd::read(handle, &mut buf)) {
5811 Ok(n) => break n,
5812 Err(e) if e.raw_os_error() == Some(libc::EINTR) => {
5813 vm.check_signals()?;
5814 continue;
5815 }
5816 Err(e) if e.raw_os_error() == Some(libc::EAGAIN) => {
5818 return Ok(None);
5819 }
5820 Err(e) => return Err(Self::io_error(zelf, e, vm)),
5821 }
5822 };
5823
5824 Ok(Some(ret))
5825 }
5826
5827 #[pymethod]
5828 fn write(
5829 zelf: &Py<Self>,
5830 obj: ArgBytesLike,
5831 vm: &VirtualMachine,
5832 ) -> PyResult<Option<usize>> {
5833 if !zelf.mode.load().contains(Mode::WRITABLE) {
5834 return Err(new_unsupported_operation(
5835 vm,
5836 "File or stream is not writable".to_owned(),
5837 ));
5838 }
5839
5840 let handle = zelf.get_fd(vm)?;
5841
5842 let len = loop {
5844 match obj.with_ref(|b| vm.allow_threads(|| crt_fd::write(handle, b))) {
5845 Ok(n) => break n,
5846 Err(e) if e.raw_os_error() == Some(libc::EINTR) => {
5847 vm.check_signals()?;
5848 continue;
5849 }
5850 Err(e) if e.raw_os_error() == Some(libc::EAGAIN) => return Ok(None),
5852 Err(e) => return Err(Self::io_error(zelf, e, vm)),
5853 }
5854 };
5855
5856 Ok(Some(len))
5858 }
5859
5860 #[pymethod]
5861 fn close(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
5862 let res = iobase_close(zelf.as_object(), vm);
5863 if !zelf.closefd.load() {
5864 zelf.fd.store(-1);
5865 return res;
5866 }
5867 let flush_exc = res.err();
5868 if zelf.finalizing.load() {
5869 Self::dealloc_warn(zelf, zelf.as_object().to_owned(), vm);
5870 }
5871 let fd = zelf.fd.swap(-1);
5872 let close_err = if fd >= 0 {
5873 crt_fd::close(unsafe { crt_fd::Owned::from_raw(fd) })
5874 .map_err(|err| Self::io_error(zelf, err, vm))
5875 .err()
5876 } else {
5877 None
5878 };
5879 match (flush_exc, close_err) {
5880 (Some(fe), Some(ce)) => {
5881 ce.set___context__(Some(fe));
5882 Err(ce)
5883 }
5884 (Some(e), None) | (None, Some(e)) => Err(e),
5885 (None, None) => Ok(()),
5886 }
5887 }
5888
5889 #[pymethod]
5890 fn seekable(&self, vm: &VirtualMachine) -> PyResult<bool> {
5891 let fd = self.get_fd(vm)?;
5892 Ok(self.seekable.load().unwrap_or_else(|| {
5893 let seekable = os::lseek(fd, 0, libc::SEEK_CUR, vm).is_ok();
5894 self.seekable.store(Some(seekable));
5895 seekable
5896 }))
5897 }
5898
5899 #[pymethod]
5900 fn seek(
5901 &self,
5902 offset: PyObjectRef,
5903 how: OptionalArg<i32>,
5904 vm: &VirtualMachine,
5905 ) -> PyResult<Offset> {
5906 let how = how.unwrap_or(0);
5907 let fd = self.get_fd(vm)?;
5908 let offset = get_offset(offset, vm)?;
5909
5910 os::lseek(fd, offset, how, vm)
5911 }
5912
5913 #[pymethod]
5914 fn tell(&self, vm: &VirtualMachine) -> PyResult<Offset> {
5915 let fd = self.get_fd(vm)?;
5916 os::lseek(fd, 0, libc::SEEK_CUR, vm)
5917 }
5918
5919 #[pymethod]
5920 fn truncate(&self, len: OptionalOption, vm: &VirtualMachine) -> PyResult<Offset> {
5921 let fd = self.get_fd(vm)?;
5922 let len = match len.flatten() {
5923 Some(l) => get_offset(l, vm)?,
5924 None => os::lseek(fd, 0, libc::SEEK_CUR, vm)?,
5925 };
5926 os::ftruncate(fd, len).map_err(|e| e.into_pyexception(vm))?;
5927 Ok(len)
5928 }
5929
5930 #[pymethod]
5931 fn isatty(&self, vm: &VirtualMachine) -> PyResult<bool> {
5932 let fd = self.fileno(vm)?;
5933 Ok(os::isatty(fd))
5934 }
5935
5936 #[pymethod]
5937 fn __getstate__(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
5938 Err(vm.new_type_error(format!("cannot pickle '{}' instances", zelf.class().name())))
5939 }
5940
5941 #[pymethod(name = "_dealloc_warn")]
5943 fn _dealloc_warn_method(
5944 zelf: &Py<Self>,
5945 source: PyObjectRef,
5946 vm: &VirtualMachine,
5947 ) -> PyResult<()> {
5948 Self::dealloc_warn(zelf, source, vm);
5949 Ok(())
5950 }
5951 }
5952
5953 impl FileIO {
5954 fn dealloc_warn(zelf: &Py<Self>, source: PyObjectRef, vm: &VirtualMachine) {
5956 if zelf.fd.load() >= 0 && zelf.closefd.load() {
5957 let repr = source
5958 .repr(vm)
5959 .map(|s| s.as_wtf8().to_owned())
5960 .unwrap_or_else(|_| Wtf8Buf::from("<file>"));
5961 if let Err(e) = crate::stdlib::_warnings::warn(
5962 vm.ctx.exceptions.resource_warning,
5963 format!("unclosed file {repr}"),
5964 1,
5965 vm,
5966 ) {
5967 vm.run_unraisable(e, None, zelf.as_object().to_owned());
5968 }
5969 }
5970 }
5971 }
5972
5973 impl Destructor for FileIO {
5974 fn slot_del(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
5975 if let Some(fileio) = zelf.downcast_ref::<FileIO>() {
5976 fileio.finalizing.store(true);
5977 }
5978 iobase_finalize(zelf, vm);
5979 Ok(())
5980 }
5981
5982 #[cold]
5983 fn del(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<()> {
5984 unreachable!("slot_del is implemented")
5985 }
5986 }
5987}
5988
5989#[cfg(all(feature = "host_env", windows))]
5991#[pymodule]
5992mod winconsoleio {
5993 use super::{_io::*, iobase_finalize};
5994 use crate::{
5995 AsObject, Py, PyObject, PyObjectRef, PyRef, PyResult, TryFromObject, VirtualMachine,
5996 builtins::{PyBaseExceptionRef, PyUtf8StrRef},
5997 common::{lock::PyMutex, wtf8::Wtf8Buf},
5998 convert::{IntoPyException, ToPyException},
5999 function::{ArgBytesLike, ArgMemoryBuffer, OptionalArg},
6000 types::{Constructor, DefaultConstructor, Destructor, Initializer, Representable},
6001 };
6002 use crossbeam_utils::atomic::AtomicCell;
6003 use windows_sys::Win32::{
6004 Foundation::{self, GENERIC_READ, GENERIC_WRITE, INVALID_HANDLE_VALUE},
6005 Globalization::{CP_UTF8, MultiByteToWideChar, WideCharToMultiByte},
6006 Storage::FileSystem::{
6007 CreateFileW, FILE_SHARE_READ, FILE_SHARE_WRITE, GetFullPathNameW, OPEN_EXISTING,
6008 },
6009 System::Console::{
6010 GetConsoleMode, GetNumberOfConsoleInputEvents, ReadConsoleW, WriteConsoleW,
6011 },
6012 };
6013
6014 type HANDLE = Foundation::HANDLE;
6015
6016 const SMALLBUF: usize = 4;
6017 const BUFMAX: usize = 32 * 1024 * 1024;
6018
6019 fn handle_from_fd(fd: i32) -> HANDLE {
6020 unsafe { rustpython_common::suppress_iph!(libc::get_osfhandle(fd)) as HANDLE }
6021 }
6022
6023 fn is_invalid_handle(handle: HANDLE) -> bool {
6024 handle == INVALID_HANDLE_VALUE || handle.is_null()
6025 }
6026
6027 fn get_console_type(handle: HANDLE) -> char {
6029 if is_invalid_handle(handle) {
6030 return '\0';
6031 }
6032 let mut mode: u32 = 0;
6033 if unsafe { GetConsoleMode(handle, &mut mode) } == 0 {
6034 return '\0';
6035 }
6036 let mut peek_count: u32 = 0;
6037 if unsafe { GetNumberOfConsoleInputEvents(handle, &mut peek_count) } != 0 {
6038 'r'
6039 } else {
6040 'w'
6041 }
6042 }
6043
6044 pub(super) fn pyio_get_console_type(path_or_fd: &PyObject, vm: &VirtualMachine) -> char {
6047 if let Ok(fd) = i32::try_from_object(vm, path_or_fd.to_owned()) {
6049 if fd >= 0 {
6050 let handle = handle_from_fd(fd);
6051 return get_console_type(handle);
6052 }
6053 return '\0';
6054 }
6055
6056 let Ok(name) = path_or_fd.str(vm) else {
6058 return '\0';
6059 };
6060 let Some(name_str) = name.to_str() else {
6061 return '\0';
6063 };
6064
6065 if name_str.eq_ignore_ascii_case("CONIN$") {
6066 return 'r';
6067 }
6068 if name_str.eq_ignore_ascii_case("CONOUT$") {
6069 return 'w';
6070 }
6071 if name_str.eq_ignore_ascii_case("CON") {
6072 return 'x';
6073 }
6074
6075 let wide: Vec<u16> = name_str.encode_utf16().chain(core::iter::once(0)).collect();
6077 let mut buf = [0u16; 260]; let length = unsafe {
6079 GetFullPathNameW(
6080 wide.as_ptr(),
6081 buf.len() as u32,
6082 buf.as_mut_ptr(),
6083 core::ptr::null_mut(),
6084 )
6085 };
6086 if length == 0 || length as usize > buf.len() {
6087 return '\0';
6088 }
6089 let full_path = &buf[..length as usize];
6090 let path_part = if full_path.len() >= 4
6092 && full_path[0] == b'\\' as u16
6093 && full_path[1] == b'\\' as u16
6094 && (full_path[2] == b'.' as u16 || full_path[2] == b'?' as u16)
6095 && full_path[3] == b'\\' as u16
6096 {
6097 &full_path[4..]
6098 } else {
6099 full_path
6100 };
6101
6102 let path_str = String::from_utf16_lossy(path_part);
6103 if path_str.eq_ignore_ascii_case("CONIN$") {
6104 'r'
6105 } else if path_str.eq_ignore_ascii_case("CONOUT$") {
6106 'w'
6107 } else if path_str.eq_ignore_ascii_case("CON") {
6108 'x'
6109 } else {
6110 '\0'
6111 }
6112 }
6113
6114 fn find_last_utf8_boundary(buf: &[u8], len: usize) -> usize {
6116 let len = len.min(buf.len());
6117 for count in 1..=4.min(len) {
6118 let c = buf[len - count];
6119 if c < 0x80 {
6120 return len;
6121 }
6122 if c >= 0xc0 {
6123 let expected = if c < 0xe0 {
6124 2
6125 } else if c < 0xf0 {
6126 3
6127 } else {
6128 4
6129 };
6130 if count < expected {
6131 return len - count;
6133 }
6134 return len;
6135 }
6136 }
6137 len
6138 }
6139
6140 #[pyattr]
6141 #[pyclass(module = "_io", name = "_WindowsConsoleIO", base = _RawIOBase)]
6142 #[derive(Debug)]
6143 pub(super) struct WindowsConsoleIO {
6144 _base: _RawIOBase,
6145 fd: AtomicCell<i32>,
6146 readable: AtomicCell<bool>,
6147 writable: AtomicCell<bool>,
6148 closefd: AtomicCell<bool>,
6149 finalizing: AtomicCell<bool>,
6150 blksize: AtomicCell<i64>,
6151 buf: PyMutex<[u8; SMALLBUF]>,
6152 }
6153
6154 impl Default for WindowsConsoleIO {
6155 fn default() -> Self {
6156 Self {
6157 _base: Default::default(),
6158 fd: AtomicCell::new(-1),
6159 readable: AtomicCell::new(false),
6160 writable: AtomicCell::new(false),
6161 closefd: AtomicCell::new(false),
6162 finalizing: AtomicCell::new(false),
6163 blksize: AtomicCell::new(super::DEFAULT_BUFFER_SIZE as _),
6164 buf: PyMutex::new([0u8; SMALLBUF]),
6165 }
6166 }
6167 }
6168
6169 impl DefaultConstructor for WindowsConsoleIO {}
6170
6171 #[derive(FromArgs)]
6172 pub struct WindowsConsoleIOArgs {
6173 #[pyarg(positional)]
6174 name: PyObjectRef,
6175 #[pyarg(any, default)]
6176 mode: Option<PyUtf8StrRef>,
6177 #[pyarg(any, default = true)]
6178 closefd: bool,
6179 #[allow(dead_code)]
6180 #[pyarg(any, default)]
6181 opener: Option<PyObjectRef>,
6182 }
6183
6184 impl Initializer for WindowsConsoleIO {
6185 type Args = WindowsConsoleIOArgs;
6186
6187 fn init(zelf: PyRef<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> {
6188 let nameobj = args.name;
6189
6190 if zelf.fd.load() >= 0 {
6191 if zelf.closefd.load() {
6192 internal_close(&zelf);
6193 } else {
6194 zelf.fd.store(-1);
6195 }
6196 }
6197
6198 if nameobj.class().is(vm.ctx.types.bool_type) {
6200 crate::stdlib::_warnings::warn(
6201 vm.ctx.exceptions.runtime_warning,
6202 "bool is used as a file descriptor".to_owned(),
6203 1,
6204 vm,
6205 )?;
6206 }
6207
6208 let mut fd: i32 = -1;
6210 if let Some(i) = nameobj.downcast_ref::<crate::builtins::PyInt>() {
6211 fd = i.try_to_primitive::<i32>(vm).unwrap_or(-1);
6212 if fd < 0 {
6213 return Err(vm.new_value_error("negative file descriptor"));
6214 }
6215 }
6216
6217 let mode_str: &str = args
6219 .mode
6220 .as_ref()
6221 .map(|s: &PyUtf8StrRef| s.as_str())
6222 .unwrap_or("r");
6223
6224 let mut rwa = false;
6225 let mut readable = false;
6226 let mut writable = false;
6227 let mut console_type = '\0';
6228 for c in mode_str.bytes() {
6229 match c {
6230 b'+' | b'a' | b'b' | b'x' => {}
6231 b'r' => {
6232 if rwa {
6233 return Err(
6234 vm.new_value_error("Must have exactly one of read or write mode")
6235 );
6236 }
6237 rwa = true;
6238 readable = true;
6239 }
6240 b'w' => {
6241 if rwa {
6242 return Err(
6243 vm.new_value_error("Must have exactly one of read or write mode")
6244 );
6245 }
6246 rwa = true;
6247 writable = true;
6248 }
6249 _ => {
6250 return Err(vm.new_value_error(format!("invalid mode: {mode_str}")));
6251 }
6252 }
6253 }
6254 if !rwa {
6255 return Err(vm.new_value_error("Must have exactly one of read or write mode"));
6256 }
6257
6258 zelf.readable.store(readable);
6259 zelf.writable.store(writable);
6260
6261 let mut _name_wide: Option<Vec<u16>> = None;
6262
6263 if fd < 0 {
6264 console_type = pyio_get_console_type(&nameobj, vm);
6266 if console_type == 'x' {
6267 if writable {
6268 console_type = 'w';
6269 } else {
6270 console_type = 'r';
6271 }
6272 }
6273
6274 zelf.closefd.store(true);
6276 if !args.closefd {
6277 return Err(vm.new_value_error("Cannot use closefd=False with file name"));
6278 }
6279
6280 let name_str = nameobj.str(vm)?;
6281 let wide: Vec<u16> = name_str
6282 .as_wtf8()
6283 .encode_wide()
6284 .chain(core::iter::once(0))
6285 .collect();
6286
6287 let access = if writable {
6288 GENERIC_WRITE
6289 } else {
6290 GENERIC_READ
6291 };
6292
6293 let mut handle: HANDLE = unsafe {
6295 CreateFileW(
6296 wide.as_ptr(),
6297 GENERIC_READ | GENERIC_WRITE,
6298 FILE_SHARE_READ | FILE_SHARE_WRITE,
6299 core::ptr::null(),
6300 OPEN_EXISTING,
6301 0,
6302 core::ptr::null_mut(),
6303 )
6304 };
6305 if is_invalid_handle(handle) {
6306 handle = unsafe {
6307 CreateFileW(
6308 wide.as_ptr(),
6309 access,
6310 FILE_SHARE_READ | FILE_SHARE_WRITE,
6311 core::ptr::null(),
6312 OPEN_EXISTING,
6313 0,
6314 core::ptr::null_mut(),
6315 )
6316 };
6317 }
6318
6319 if is_invalid_handle(handle) {
6320 return Err(std::io::Error::last_os_error().to_pyexception(vm));
6321 }
6322
6323 let osf_flags = if writable {
6324 libc::O_WRONLY | libc::O_BINARY | 0x80 } else {
6326 libc::O_RDONLY | libc::O_BINARY | 0x80 };
6328
6329 fd = unsafe { libc::open_osfhandle(handle as isize, osf_flags) };
6330 if fd < 0 {
6331 unsafe {
6332 Foundation::CloseHandle(handle);
6333 }
6334 return Err(std::io::Error::last_os_error().to_pyexception(vm));
6335 }
6336
6337 _name_wide = Some(wide);
6338 } else {
6339 zelf.closefd.store(false);
6341 }
6342
6343 zelf.fd.store(fd);
6344
6345 if console_type == '\0' {
6347 let handle = handle_from_fd(fd);
6348 console_type = get_console_type(handle);
6349 }
6350
6351 if console_type == '\0' {
6352 internal_close(&zelf);
6354 return Err(vm.new_value_error("Cannot open non-console file"));
6355 }
6356
6357 if writable && console_type != 'w' {
6358 internal_close(&zelf);
6359 return Err(vm.new_value_error("Cannot open console input buffer for writing"));
6360 }
6361 if readable && console_type != 'r' {
6362 internal_close(&zelf);
6363 return Err(vm.new_value_error("Cannot open console output buffer for reading"));
6364 }
6365
6366 zelf.blksize.store(super::DEFAULT_BUFFER_SIZE as _);
6367 *zelf.buf.lock() = [0u8; SMALLBUF];
6368
6369 zelf.as_object().set_attr("name", nameobj, vm)?;
6370
6371 Ok(())
6372 }
6373 }
6374
6375 fn internal_close(zelf: &WindowsConsoleIO) {
6376 let fd = zelf.fd.swap(-1);
6377 if fd >= 0 && zelf.closefd.load() {
6378 unsafe {
6379 libc::close(fd);
6380 }
6381 }
6382 }
6383
6384 impl Representable for WindowsConsoleIO {
6385 #[inline]
6386 fn repr_str(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
6387 let type_name = zelf.class().slot_name();
6388 let fd = zelf.fd.load();
6389 if fd < 0 {
6390 return Ok(format!("<{type_name} [closed]>"));
6391 }
6392 let mode = if zelf.readable.load() { "rb" } else { "wb" };
6393 let closefd = if zelf.closefd.load() { "True" } else { "False" };
6394 Ok(format!("<{type_name} mode='{mode}' closefd={closefd}>"))
6395 }
6396 }
6397
6398 #[pyclass(
6399 with(Constructor, Initializer, Representable, Destructor),
6400 flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
6401 )]
6402 impl WindowsConsoleIO {
6403 #[allow(dead_code)]
6404 fn io_error(
6405 zelf: &Py<Self>,
6406 error: std::io::Error,
6407 vm: &VirtualMachine,
6408 ) -> PyBaseExceptionRef {
6409 let exc = error.to_pyexception(vm);
6410 if let Ok(name) = zelf.as_object().get_attr("name", vm) {
6411 exc.as_object()
6412 .set_attr("filename", name, vm)
6413 .expect("OSError.filename set must succeed");
6414 }
6415 exc
6416 }
6417
6418 #[pygetset]
6419 fn closed(&self) -> bool {
6420 self.fd.load() < 0
6421 }
6422
6423 #[pygetset]
6424 fn closefd(&self) -> bool {
6425 self.closefd.load()
6426 }
6427
6428 #[pygetset(name = "_blksize")]
6429 fn blksize(&self) -> i64 {
6430 self.blksize.load()
6431 }
6432
6433 #[pygetset]
6434 fn mode(&self) -> &'static str {
6435 if self.readable.load() { "rb" } else { "wb" }
6436 }
6437
6438 #[pymethod]
6439 fn fileno(&self, vm: &VirtualMachine) -> PyResult<i32> {
6440 let fd = self.fd.load();
6441 if fd >= 0 {
6442 Ok(fd)
6443 } else {
6444 Err(io_closed_error(vm))
6445 }
6446 }
6447
6448 fn get_fd(&self, vm: &VirtualMachine) -> PyResult<i32> {
6449 self.fileno(vm)
6450 }
6451
6452 #[pymethod]
6453 fn readable(&self, vm: &VirtualMachine) -> PyResult<bool> {
6454 if self.fd.load() < 0 {
6455 return Err(io_closed_error(vm));
6456 }
6457 Ok(self.readable.load())
6458 }
6459
6460 #[pymethod]
6461 fn writable(&self, vm: &VirtualMachine) -> PyResult<bool> {
6462 if self.fd.load() < 0 {
6463 return Err(io_closed_error(vm));
6464 }
6465 Ok(self.writable.load())
6466 }
6467
6468 #[pymethod]
6469 fn isatty(&self, vm: &VirtualMachine) -> PyResult<bool> {
6470 if self.fd.load() < 0 {
6471 return Err(io_closed_error(vm));
6472 }
6473 Ok(true)
6474 }
6475
6476 #[pymethod]
6477 fn close(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
6478 let res = iobase_close(zelf.as_object(), vm);
6479 if !zelf.closefd.load() {
6480 zelf.fd.store(-1);
6481 return res;
6482 }
6483 let flush_exc = res.err();
6484 if zelf.finalizing.load() {
6485 Self::dealloc_warn(zelf, zelf.as_object().to_owned(), vm);
6486 }
6487 let fd = zelf.fd.swap(-1);
6488 let close_err: Option<PyBaseExceptionRef> = if fd >= 0 {
6489 let result = unsafe { libc::close(fd) };
6490 if result < 0 {
6491 Some(std::io::Error::last_os_error().into_pyexception(vm))
6492 } else {
6493 None
6494 }
6495 } else {
6496 None
6497 };
6498 match (flush_exc, close_err) {
6499 (Some(fe), Some(ce)) => {
6500 ce.set___context__(Some(fe));
6501 Err(ce)
6502 }
6503 (Some(e), None) | (None, Some(e)) => Err(e),
6504 (None, None) => Ok(()),
6505 }
6506 }
6507
6508 fn dealloc_warn(zelf: &Py<Self>, source: PyObjectRef, vm: &VirtualMachine) {
6509 if zelf.fd.load() >= 0 && zelf.closefd.load() {
6510 let repr = source
6511 .repr(vm)
6512 .map(|s| s.as_wtf8().to_owned())
6513 .unwrap_or_else(|_| Wtf8Buf::from("<file>"));
6514 if let Err(e) = crate::stdlib::_warnings::warn(
6515 vm.ctx.exceptions.resource_warning,
6516 format!("unclosed file {repr}"),
6517 1,
6518 vm,
6519 ) {
6520 vm.run_unraisable(e, None, zelf.as_object().to_owned());
6521 }
6522 }
6523 }
6524
6525 fn copy_from_buf(buf: &mut [u8; SMALLBUF], dest: &mut [u8]) -> usize {
6526 let mut n = 0;
6527 while buf[0] != 0 && n < dest.len() {
6528 dest[n] = buf[0];
6529 n += 1;
6530 for i in 1..SMALLBUF {
6531 buf[i - 1] = buf[i];
6532 }
6533 buf[SMALLBUF - 1] = 0;
6534 }
6535 n
6536 }
6537
6538 #[pymethod]
6539 fn readinto(&self, buffer: ArgMemoryBuffer, vm: &VirtualMachine) -> PyResult<usize> {
6540 let fd = self.get_fd(vm)?;
6541 if !self.readable.load() {
6542 return Err(new_unsupported_operation(
6543 vm,
6544 "Console buffer does not support reading".to_owned(),
6545 ));
6546 }
6547 let mut buf_ref = buffer.borrow_buf_mut();
6548 let len = buf_ref.len();
6549 if len == 0 {
6550 return Ok(0);
6551 }
6552 if len > BUFMAX {
6553 return Err(vm.new_value_error(format!("cannot read more than {BUFMAX} bytes")));
6554 }
6555
6556 let handle = handle_from_fd(fd);
6557 if is_invalid_handle(handle) {
6558 return Err(std::io::Error::last_os_error().to_pyexception(vm));
6559 }
6560
6561 let mut wlen = (len / 4) as u32;
6563 if wlen == 0 {
6564 wlen = 1;
6565 }
6566
6567 let dest = &mut *buf_ref;
6568
6569 let mut read_len = {
6571 let mut buf = self.buf.lock();
6572 Self::copy_from_buf(&mut buf, dest)
6573 };
6574 if read_len > 0 {
6575 wlen = wlen.saturating_sub(1);
6576 }
6577 if read_len >= len || wlen == 0 {
6578 return Ok(read_len);
6579 }
6580
6581 let mut wbuf = vec![0u16; wlen as usize];
6583 let mut nread: u32 = 0;
6584 let res = unsafe {
6585 ReadConsoleW(
6586 handle,
6587 wbuf.as_mut_ptr() as _,
6588 wlen,
6589 &mut nread,
6590 core::ptr::null(),
6591 )
6592 };
6593 if res == 0 {
6594 return Err(std::io::Error::last_os_error().into_pyexception(vm));
6595 }
6596 if nread == 0 {
6597 return Ok(read_len);
6598 }
6599
6600 if nread > 0 && wbuf[0] == 0x1A {
6602 return Ok(read_len);
6603 }
6604
6605 let remaining = len - read_len;
6607 let u8n;
6608 if remaining < 4 {
6609 let mut buf = self.buf.lock();
6611 let converted = unsafe {
6612 WideCharToMultiByte(
6613 CP_UTF8,
6614 0,
6615 wbuf.as_ptr(),
6616 nread as i32,
6617 buf.as_mut_ptr() as _,
6618 SMALLBUF as i32,
6619 core::ptr::null(),
6620 core::ptr::null_mut(),
6621 )
6622 };
6623 if converted > 0 {
6624 u8n = Self::copy_from_buf(&mut buf, &mut dest[read_len..]) as i32;
6625 } else {
6626 u8n = 0;
6627 }
6628 } else {
6629 u8n = unsafe {
6630 WideCharToMultiByte(
6631 CP_UTF8,
6632 0,
6633 wbuf.as_ptr(),
6634 nread as i32,
6635 dest[read_len..].as_mut_ptr() as _,
6636 remaining as i32,
6637 core::ptr::null(),
6638 core::ptr::null_mut(),
6639 )
6640 };
6641 }
6642
6643 if u8n > 0 {
6644 read_len += u8n as usize;
6645 } else {
6646 let err = std::io::Error::last_os_error();
6647 if err.raw_os_error() == Some(122) {
6648 let needed = unsafe {
6650 WideCharToMultiByte(
6651 CP_UTF8,
6652 0,
6653 wbuf.as_ptr(),
6654 nread as i32,
6655 core::ptr::null_mut(),
6656 0,
6657 core::ptr::null(),
6658 core::ptr::null_mut(),
6659 )
6660 };
6661 if needed > 0 {
6662 return Err(vm.new_system_error(format!(
6663 "Buffer had room for {remaining} bytes but {needed} bytes required",
6664 )));
6665 }
6666 }
6667 return Err(err.into_pyexception(vm));
6668 }
6669
6670 Ok(read_len)
6671 }
6672
6673 #[pymethod]
6674 fn readall(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
6675 if self.fd.load() < 0 {
6676 return Err(io_closed_error(vm));
6677 }
6678
6679 let handle = handle_from_fd(self.fd.load());
6680 if is_invalid_handle(handle) {
6681 return Err(std::io::Error::last_os_error().to_pyexception(vm));
6682 }
6683
6684 let mut result = Vec::new();
6685
6686 {
6688 let mut buf = self.buf.lock();
6689 let mut tmp = [0u8; SMALLBUF];
6690 let n = Self::copy_from_buf(&mut buf, &mut tmp);
6691 result.extend_from_slice(&tmp[..n]);
6692 }
6693
6694 let mut wbuf = vec![0u16; 8192];
6695 loop {
6696 let mut nread: u32 = 0;
6697 let res = unsafe {
6698 ReadConsoleW(
6699 handle,
6700 wbuf.as_mut_ptr() as _,
6701 wbuf.len() as u32,
6702 &mut nread,
6703 core::ptr::null(),
6704 )
6705 };
6706 if res == 0 {
6707 return Err(std::io::Error::last_os_error().into_pyexception(vm));
6708 }
6709 if nread == 0 {
6710 break;
6711 }
6712 if wbuf[0] == 0x1A {
6714 break;
6715 }
6716 let needed = unsafe {
6718 WideCharToMultiByte(
6719 CP_UTF8,
6720 0,
6721 wbuf.as_ptr(),
6722 nread as i32,
6723 core::ptr::null_mut(),
6724 0,
6725 core::ptr::null(),
6726 core::ptr::null_mut(),
6727 )
6728 };
6729 if needed == 0 {
6730 return Err(std::io::Error::last_os_error().into_pyexception(vm));
6731 }
6732 let offset = result.len();
6733 result.resize(offset + needed as usize, 0);
6734 let written = unsafe {
6735 WideCharToMultiByte(
6736 CP_UTF8,
6737 0,
6738 wbuf.as_ptr(),
6739 nread as i32,
6740 result[offset..].as_mut_ptr() as _,
6741 needed,
6742 core::ptr::null(),
6743 core::ptr::null_mut(),
6744 )
6745 };
6746 if written == 0 {
6747 return Err(std::io::Error::last_os_error().into_pyexception(vm));
6748 }
6749 if nread < wbuf.len() as u32 {
6751 break;
6752 }
6753 }
6754
6755 Ok(vm.ctx.new_bytes(result).into())
6756 }
6757
6758 #[pymethod]
6759 fn read(&self, size: OptionalArg<isize>, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
6760 if self.fd.load() < 0 {
6761 return Err(io_closed_error(vm));
6762 }
6763 if !self.readable.load() {
6764 return Err(new_unsupported_operation(
6765 vm,
6766 "Console buffer does not support reading".to_owned(),
6767 ));
6768 }
6769 let size = size.unwrap_or(-1);
6770 if size < 0 {
6771 return self.readall(vm);
6772 }
6773 if size as usize > BUFMAX {
6774 return Err(vm.new_value_error(format!("cannot read more than {BUFMAX} bytes")));
6775 }
6776 let mut buf = vec![0u8; size as usize];
6777 let handle = handle_from_fd(self.fd.load());
6778 if is_invalid_handle(handle) {
6779 return Err(std::io::Error::last_os_error().to_pyexception(vm));
6780 }
6781
6782 let len = size as usize;
6783
6784 let mut wlen = (len / 4) as u32;
6785 if wlen == 0 {
6786 wlen = 1;
6787 }
6788
6789 let mut read_len = {
6790 let mut ibuf = self.buf.lock();
6791 Self::copy_from_buf(&mut ibuf, &mut buf)
6792 };
6793 if read_len > 0 {
6794 wlen = wlen.saturating_sub(1);
6795 }
6796 if read_len >= len || wlen == 0 {
6797 buf.truncate(read_len);
6798 return Ok(vm.ctx.new_bytes(buf).into());
6799 }
6800
6801 let mut wbuf = vec![0u16; wlen as usize];
6802 let mut nread: u32 = 0;
6803 let res = unsafe {
6804 ReadConsoleW(
6805 handle,
6806 wbuf.as_mut_ptr() as _,
6807 wlen,
6808 &mut nread,
6809 core::ptr::null(),
6810 )
6811 };
6812 if res == 0 {
6813 return Err(std::io::Error::last_os_error().into_pyexception(vm));
6814 }
6815 if nread == 0 || wbuf[0] == 0x1A {
6816 buf.truncate(read_len);
6817 return Ok(vm.ctx.new_bytes(buf).into());
6818 }
6819
6820 let remaining = len - read_len;
6821 let u8n;
6822 if remaining < 4 {
6823 let mut ibuf = self.buf.lock();
6824 let converted = unsafe {
6825 WideCharToMultiByte(
6826 CP_UTF8,
6827 0,
6828 wbuf.as_ptr(),
6829 nread as i32,
6830 ibuf.as_mut_ptr() as _,
6831 SMALLBUF as i32,
6832 core::ptr::null(),
6833 core::ptr::null_mut(),
6834 )
6835 };
6836 if converted > 0 {
6837 u8n = Self::copy_from_buf(&mut ibuf, &mut buf[read_len..]) as i32;
6838 } else {
6839 u8n = 0;
6840 }
6841 } else {
6842 u8n = unsafe {
6843 WideCharToMultiByte(
6844 CP_UTF8,
6845 0,
6846 wbuf.as_ptr(),
6847 nread as i32,
6848 buf[read_len..].as_mut_ptr() as _,
6849 remaining as i32,
6850 core::ptr::null(),
6851 core::ptr::null_mut(),
6852 )
6853 };
6854 }
6855
6856 if u8n > 0 {
6857 read_len += u8n as usize;
6858 } else {
6859 let err = std::io::Error::last_os_error();
6860 if err.raw_os_error() == Some(122) {
6861 let needed = unsafe {
6863 WideCharToMultiByte(
6864 CP_UTF8,
6865 0,
6866 wbuf.as_ptr(),
6867 nread as i32,
6868 core::ptr::null_mut(),
6869 0,
6870 core::ptr::null(),
6871 core::ptr::null_mut(),
6872 )
6873 };
6874 if needed > 0 {
6875 return Err(vm.new_system_error(format!(
6876 "Buffer had room for {remaining} bytes but {needed} bytes required",
6877 )));
6878 }
6879 }
6880 return Err(err.into_pyexception(vm));
6881 }
6882
6883 buf.truncate(read_len);
6884 Ok(vm.ctx.new_bytes(buf).into())
6885 }
6886
6887 #[pymethod]
6888 fn write(&self, b: ArgBytesLike, vm: &VirtualMachine) -> PyResult<usize> {
6889 if self.fd.load() < 0 {
6890 return Err(io_closed_error(vm));
6891 }
6892 if !self.writable.load() {
6893 return Err(new_unsupported_operation(
6894 vm,
6895 "Console buffer does not support writing".to_owned(),
6896 ));
6897 }
6898
6899 let handle = handle_from_fd(self.fd.load());
6900 if is_invalid_handle(handle) {
6901 return Err(std::io::Error::last_os_error().to_pyexception(vm));
6902 }
6903
6904 let data = b.borrow_buf();
6905 let data = &*data;
6906 if data.is_empty() {
6907 return Ok(0);
6908 }
6909
6910 let mut len = data.len().min(BUFMAX);
6911
6912 let max_wlen: u32 = 32766 / 2;
6914 len = len.min(max_wlen as usize * 3);
6915
6916 let wlen;
6918 loop {
6919 len = find_last_utf8_boundary(data, len);
6920 let w = unsafe {
6921 MultiByteToWideChar(
6922 CP_UTF8,
6923 0,
6924 data.as_ptr(),
6925 len as i32,
6926 core::ptr::null_mut(),
6927 0,
6928 )
6929 };
6930 if w as u32 <= max_wlen {
6931 wlen = w;
6932 break;
6933 }
6934 len /= 2;
6935 }
6936 if wlen == 0 {
6937 return Ok(0);
6938 }
6939
6940 let mut wbuf = vec![0u16; wlen as usize];
6941 let wlen = unsafe {
6942 MultiByteToWideChar(
6943 CP_UTF8,
6944 0,
6945 data.as_ptr(),
6946 len as i32,
6947 wbuf.as_mut_ptr(),
6948 wlen,
6949 )
6950 };
6951 if wlen == 0 {
6952 return Err(std::io::Error::last_os_error().into_pyexception(vm));
6953 }
6954
6955 let mut n_written: u32 = 0;
6956 let res = unsafe {
6957 WriteConsoleW(
6958 handle,
6959 wbuf.as_ptr() as _,
6960 wlen as u32,
6961 &mut n_written,
6962 core::ptr::null(),
6963 )
6964 };
6965 if res == 0 {
6966 return Err(std::io::Error::last_os_error().into_pyexception(vm));
6967 }
6968
6969 if n_written < wlen as u32 {
6971 len = wchar_to_utf8_count(data, len, n_written);
6973 }
6974
6975 Ok(len)
6976 }
6977
6978 #[pymethod(name = "__reduce__")]
6979 fn reduce(_zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult {
6980 Err(vm.new_type_error("cannot pickle '_WindowsConsoleIO' instances"))
6981 }
6982 }
6983
6984 fn wchar_to_utf8_count(data: &[u8], mut len: usize, mut n: u32) -> usize {
6986 let mut start: usize = 0;
6987 loop {
6988 let mut mid = 0;
6989 for i in (len / 2)..=len {
6990 mid = find_last_utf8_boundary(data, i);
6991 if mid != 0 {
6992 break;
6993 }
6994 }
6995 if mid == len {
6996 return start + len;
6997 }
6998 if mid == 0 {
6999 mid = if len > 1 { len - 1 } else { 1 };
7000 }
7001 let wlen = unsafe {
7002 MultiByteToWideChar(
7003 CP_UTF8,
7004 0,
7005 data[start..].as_ptr(),
7006 mid as i32,
7007 core::ptr::null_mut(),
7008 0,
7009 )
7010 } as u32;
7011 if wlen <= n {
7012 start += mid;
7013 len -= mid;
7014 n -= wlen;
7015 } else {
7016 len = mid;
7017 }
7018 }
7019 }
7020
7021 impl Destructor for WindowsConsoleIO {
7022 fn slot_del(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
7023 if let Some(cio) = zelf.downcast_ref::<WindowsConsoleIO>() {
7024 cio.finalizing.store(true);
7025 }
7026 iobase_finalize(zelf, vm);
7027 Ok(())
7028 }
7029
7030 #[cold]
7031 fn del(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<()> {
7032 unreachable!("slot_del is implemented")
7033 }
7034 }
7035}