@@ -41,24 +41,26 @@ namespace fs = std::filesystem;
4141namespace py {
4242
4343namespace {
44- static PyType *s_io_base = nullptr ;
45- static PyType *s_io_raw_iobase = nullptr ;
46- static PyType *s_io_buffered_io_base = nullptr ;
47- static PyType *s_io_buffered_reader = nullptr ;
48- static PyType *s_io_buffered_writer = nullptr ;
49- static PyType *s_io_buffered_rwpair = nullptr ;
50- static PyType *s_io_buffered_random = nullptr ;
51- static PyType *s_io_textiobase = nullptr ;
52- static PyType *s_io_incremental_newline_decoder = nullptr ;
53- static PyType *s_io_bytesio = nullptr ;
54- static PyType *s_io_fileio = nullptr ;
55- static PyType *s_io_stringio = nullptr ;
56- static PyType *s_io_textiowrapper = nullptr ;
57-
58- static PyType *s_blocking_io_error = nullptr ;
59- static PyType *s_unsupported_operation_type = nullptr ;
44+ PyType *s_io_base = nullptr ;
45+ PyType *s_io_raw_iobase = nullptr ;
46+ PyType *s_io_buffered_io_base = nullptr ;
47+ PyType *s_io_buffered_reader = nullptr ;
48+ PyType *s_io_buffered_writer = nullptr ;
49+ PyType *s_io_buffered_rwpair = nullptr ;
50+ PyType *s_io_buffered_random = nullptr ;
51+ PyType *s_io_textiobase = nullptr ;
52+ PyType *s_io_incremental_newline_decoder = nullptr ;
53+ PyType *s_io_bytesio = nullptr ;
54+ PyType *s_io_fileio = nullptr ;
55+ PyType *s_io_stringio = nullptr ;
56+ PyType *s_io_textiowrapper = nullptr ;
57+
58+ PyType *s_blocking_io_error = nullptr ;
59+ PyType *s_unsupported_operation_type = nullptr ;
6060}// namespace
6161
62+ static constexpr size_t s_default_buffer_size = 8192 ;
63+
6264Exception *unsupported_operation (PyTuple *args, PyDict *kwargs)
6365{
6466 ASSERT (s_unsupported_operation_type);
@@ -814,7 +816,7 @@ struct Buffered
814816
815817 PyResult<std::monostate> check_initialized () const
816818 {
817- if (!ok) {
819+ if (!ok || !raw ) {
818820 if (detached) {
819821 return Err (value_error (" raw stream has been detached" ));
820822 } else {
@@ -1029,7 +1031,8 @@ class BufferedReader
10291031 this ->readable_ = true ;
10301032 this ->writable_ = false ;
10311033 this ->fast_closed_checks = false ;
1032- this ->ok = true ;
1034+ // the object only becomes usable once __init__ provides a raw stream
1035+ this ->ok = raw != nullptr ;
10331036 }
10341037
10351038 BufferedReader (PyType *type) : BufferedReader(type, nullptr , 0 ) {}
@@ -1359,11 +1362,13 @@ class BufferedWriter
13591362 BufferedWriter (PyType *type, PyObject *raw, int buffer_size) : BufferedIOBase(type)
13601363 {
13611364 this ->raw = raw;
1362- ( void ) buffer_size;
1365+ if (buffer_size > 0 ) { m_buffer_size = static_cast < size_t >( buffer_size); }
13631366 this ->readable_ = false ;
13641367 this ->writable_ = true ;
13651368 this ->fast_closed_checks = false ;
1366- this ->ok = true ;
1369+ // the object only becomes usable once __init__ provides a raw stream
1370+ this ->ok = raw != nullptr ;
1371+ this ->buffer = std::make_unique<std::stringbuf>();
13671372 }
13681373
13691374 BufferedWriter (PyType *type) : BufferedWriter(type, nullptr , 0 ) {}
@@ -1397,7 +1402,7 @@ class BufferedWriter
13971402 static PyResult<BufferedWriter *> create (PyObject *raw, int buffer_size)
13981403 {
13991404 auto &heap = VirtualMachine::the ().heap ();
1400- if (auto *obj = heap.allocate <BufferedWriter>(s_io_buffered_reader , raw, buffer_size)) {
1405+ if (auto *obj = heap.allocate <BufferedWriter>(s_io_buffered_writer , raw, buffer_size)) {
14011406 return Ok (obj);
14021407 }
14031408 return Err (memory_error (sizeof (BufferedWriter)));
@@ -1410,21 +1415,24 @@ class BufferedWriter
14101415
14111416 PyResult<int32_t > __init__ (PyTuple *args, PyDict *kwargs)
14121417 {
1413- ASSERT (!kwargs || kwargs->map ().empty ());
1414- if (!args || args->elements ().empty ()) {
1415- return Err (type_error (" missing required argument 'raw'" ));
1416- }
1417- if (args->elements ().size () > 1 ) { TODO (); }
1418- return PyObject::from (args->elements ()[0 ])
1419- .and_then ([this ](PyObject *raw) -> PyResult<int32_t > {
1420- this ->raw = raw;
1421- this ->readable_ = false ;
1422- this ->writable_ = true ;
1423- this ->fast_closed_checks = false ;
1424- this ->ok = true ;
1425- this ->buffer = std::make_unique<std::stringbuf>();
1426- return Ok (0 );
1427- });
1418+ auto parse_result = PyArgsParser<PyObject *, int64_t >::unpack_tuple (args,
1419+ kwargs,
1420+ " BufferedWriter" ,
1421+ std::integral_constant<size_t , 1 >{},
1422+ std::integral_constant<size_t , 2 >{},
1423+ static_cast <int64_t >(s_default_buffer_size));
1424+ if (parse_result.is_err ()) { return Err (parse_result.unwrap_err ()); }
1425+ auto [raw, buffer_size] = parse_result.unwrap ();
1426+ if (buffer_size <= 0 ) { return Err (value_error (" buffer size must be strictly positive" )); }
1427+ this ->raw = raw;
1428+ this ->readable_ = false ;
1429+ this ->writable_ = true ;
1430+ this ->fast_closed_checks = false ;
1431+ this ->ok = true ;
1432+ this ->buffer = std::make_unique<std::stringbuf>();
1433+ m_buffer_size = static_cast <size_t >(buffer_size);
1434+ m_buffered_bytes = 0 ;
1435+ return Ok (0 );
14281436 }
14291437
14301438 PyResult<PyObject *> __repr__ () const
@@ -1453,8 +1461,7 @@ class BufferedWriter
14531461 return raw->get_method (PyString::create (" write" ).unwrap ())
14541462 .and_then ([memobj](PyObject *write) {
14551463 return write->call (PyTuple::create (memobj.unwrap ()).unwrap (), nullptr );
1456- })
1457- .and_then ([](auto ) { return Ok (py_none ()); });
1464+ });
14581465 }
14591466
14601467 PyResult<PyObject *> write (PyTuple *args, PyDict *kwargs)
@@ -1473,41 +1480,66 @@ class BufferedWriter
14731480 " write() argument must be contiguous buffer, not {}" , arg->type ()->name ()));
14741481 }
14751482
1476- // implementation
1483+ if (auto err = check_initialized (); err.is_err ()) { return Err (err.unwrap_err ()); }
1484+
14771485 if (Buffered<BufferedWriter>::is_closed ()) {
14781486 return Err (value_error (" write to closed file" ));
14791487 }
14801488
1481- auto written =
1482- this ->buffer ->sputn (static_cast <char *>(buffer.buf ->get_buffer ()), buffer.len );
1483- if (written == buffer.len ) {
1484- // fast path: everything was written - we are done
1485- return PyInteger::create (written);
1489+ const char *start = static_cast <char *>(buffer.buf ->get_buffer ());
1490+ const auto len = static_cast <size_t >(buffer.len );
1491+
1492+ // fast path: the data fits in the buffer's free space, so it is only buffered and does
1493+ // not reach the raw stream
1494+ if (m_buffered_bytes + len <= m_buffer_size) {
1495+ const auto written = this ->buffer ->sputn (start, len);
1496+ ASSERT (static_cast <size_t >(written) == len);
1497+ m_buffered_bytes += len;
1498+ return PyInteger::create (len);
14861499 }
14871500
1488- // flush our internal buffer to raw
1501+ // the data does not fit: drain our buffer to raw first
14891502 if (auto r = flush (); r.is_err ()) { return r; }
14901503
1491- // TODO: this currently writes to internal buffer and then immediately makes a copy in
1492- // write_raw. This is because streambuf doesn't tell us how much capacity it has left, so we
1493- // just give it as much as we can each time
1494- ssize_t remaining = buffer.len ;
1495- char *start = static_cast <char *>(buffer.buf ->get_buffer ());
1496- // TODO: use a buffer that actually tells us how much we can write to it, rather than
1497- // hardcode 4096
1498- while (remaining > 4096 ) {
1499- // flush our internal buffer to raw
1504+ // write oversized data directly to raw, bypassing the buffer
1505+ size_t written = 0 ;
1506+ size_t remaining = len;
1507+ while (remaining > m_buffer_size) {
15001508 auto r = raw_write (start + written, remaining);
15011509 if (r.is_err ()) { return r; }
1502- ASSERT (as<PyInteger>(r.unwrap ()));
1503- const auto count = as<PyInteger>(r.unwrap ())->as_size_t ();
1504- written += count;
1505- remaining -= count;
1510+ if (r.unwrap () == py_none ()) {
1511+ // TODO: raise BlockingIOError
1512+ // a non-blocking raw stream accepted no data
1513+ return Err (os_error (" write could not complete without blocking" ));
1514+ }
1515+ if (!as<PyInteger>(r.unwrap ())) {
1516+ return Err (type_error (
1517+ " '{}' object cannot be interpreted as an integer" , r.unwrap ()->type ()->name ()));
1518+ }
1519+ const auto &count = as<PyInteger>(r.unwrap ())->as_big_int ();
1520+ if (count < 0 || count > remaining) {
1521+ return Err (
1522+ os_error (" raw write() returned invalid length {} (should have been "
1523+ " between 0 and {})" ,
1524+ count.get_str (),
1525+ remaining));
1526+ }
1527+ if (count == 0 ) {
1528+ // TODO: retry with a zero-byte raw write until a signal interrupts
1529+ // it. We have no signal check in this loop yet
1530+ return Err (os_error (" write could not complete without blocking" ));
1531+ }
1532+ written += count.get_ui ();
1533+ remaining -= count.get_ui ();
15061534 }
15071535
1508- // put the rest in our buffer. Assumes that we can write 4096 bytes
1509- const auto count = this ->buffer ->sputn (start + written, remaining);
1510- ASSERT (remaining - count == 0 );
1536+ // buffer the tail, which is now guaranteed to fit
1537+ if (remaining > 0 ) {
1538+ const auto count = this ->buffer ->sputn (start + written, remaining);
1539+ ASSERT (static_cast <size_t >(count) == remaining);
1540+ m_buffered_bytes += remaining;
1541+ written += remaining;
1542+ }
15111543
15121544 return PyInteger::create (written);
15131545 }
@@ -1538,6 +1570,10 @@ class BufferedWriter
15381570 PyObject::visit_graph (visitor);
15391571 visit_graph_buffered (visitor);
15401572 }
1573+
1574+ private:
1575+ size_t m_buffer_size{ s_default_buffer_size };
1576+ size_t m_buffered_bytes{ 0 };
15411577};
15421578
15431579template <> PyResult<PyObject *> Buffered<BufferedWriter>::flush_and_rewind()
@@ -1548,6 +1584,7 @@ template<> PyResult<PyObject *> Buffered<BufferedWriter>::flush_and_rewind()
15481584
15491585PyResult<PyObject *> BufferedWriter::flush ()
15501586{
1587+ if (auto err = check_initialized (); err.is_err ()) { return Err (err.unwrap_err ()); }
15511588 if (!writable_) { return Ok (py_none ()); }
15521589 std::array<char , BUFSIZ > to_write;
15531590
@@ -1568,6 +1605,7 @@ PyResult<PyObject *> BufferedWriter::flush()
15681605 return r;
15691606 }
15701607 }
1608+ m_buffered_bytes = 0 ;
15711609 return Ok (py_none ());
15721610}
15731611
@@ -3462,9 +3500,8 @@ PyModule *io_module()
34623500 return open (arg0, " rb" );
34633501 }));
34643502
3465- // C++ standard streams currently do not provide an API to get default buffer size, and C's
3466- // BUFSIZE doesn't have to be respected
3467- s_io_module->add_symbol (PyString::create (" DEFAULT_BUFFER_SIZE" ).unwrap (), Number{ 0 });
3503+ s_io_module->add_symbol (PyString::create (" DEFAULT_BUFFER_SIZE" ).unwrap (),
3504+ Number{ static_cast <int64_t >(s_default_buffer_size) });
34683505
34693506 // >> type("UnsupportedOperation", (_io.OSError, ValueError), {})
34703507 auto unsupported_operation_type = types::type ()->call (
0 commit comments