@@ -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);
@@ -1359,11 +1361,12 @@ class BufferedWriter
13591361 BufferedWriter (PyType *type, PyObject *raw, int buffer_size) : BufferedIOBase(type)
13601362 {
13611363 this ->raw = raw;
1362- ( void ) buffer_size;
1364+ if (buffer_size > 0 ) { m_buffer_size = static_cast < size_t >( buffer_size); }
13631365 this ->readable_ = false ;
13641366 this ->writable_ = true ;
13651367 this ->fast_closed_checks = false ;
13661368 this ->ok = true ;
1369+ this ->buffer = std::make_unique<std::stringbuf>();
13671370 }
13681371
13691372 BufferedWriter (PyType *type) : BufferedWriter(type, nullptr , 0 ) {}
@@ -1397,7 +1400,7 @@ class BufferedWriter
13971400 static PyResult<BufferedWriter *> create (PyObject *raw, int buffer_size)
13981401 {
13991402 auto &heap = VirtualMachine::the ().heap ();
1400- if (auto *obj = heap.allocate <BufferedWriter>(s_io_buffered_reader , raw, buffer_size)) {
1403+ if (auto *obj = heap.allocate <BufferedWriter>(s_io_buffered_writer , raw, buffer_size)) {
14011404 return Ok (obj);
14021405 }
14031406 return Err (memory_error (sizeof (BufferedWriter)));
@@ -1410,21 +1413,24 @@ class BufferedWriter
14101413
14111414 PyResult<int32_t > __init__ (PyTuple *args, PyDict *kwargs)
14121415 {
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- });
1416+ auto parse_result = PyArgsParser<PyObject *, int64_t >::unpack_tuple (args,
1417+ kwargs,
1418+ " BufferedWriter" ,
1419+ std::integral_constant<size_t , 1 >{},
1420+ std::integral_constant<size_t , 2 >{},
1421+ static_cast <int64_t >(s_default_buffer_size));
1422+ if (parse_result.is_err ()) { return Err (parse_result.unwrap_err ()); }
1423+ auto [raw, buffer_size] = parse_result.unwrap ();
1424+ if (buffer_size <= 0 ) { return Err (value_error (" buffer size must be strictly positive" )); }
1425+ this ->raw = raw;
1426+ this ->readable_ = false ;
1427+ this ->writable_ = true ;
1428+ this ->fast_closed_checks = false ;
1429+ this ->ok = true ;
1430+ this ->buffer = std::make_unique<std::stringbuf>();
1431+ m_buffer_size = static_cast <size_t >(buffer_size);
1432+ m_buffered_bytes = 0 ;
1433+ return Ok (0 );
14281434 }
14291435
14301436 PyResult<PyObject *> __repr__ () const
@@ -1453,8 +1459,7 @@ class BufferedWriter
14531459 return raw->get_method (PyString::create (" write" ).unwrap ())
14541460 .and_then ([memobj](PyObject *write) {
14551461 return write->call (PyTuple::create (memobj.unwrap ()).unwrap (), nullptr );
1456- })
1457- .and_then ([](auto ) { return Ok (py_none ()); });
1462+ });
14581463 }
14591464
14601465 PyResult<PyObject *> write (PyTuple *args, PyDict *kwargs)
@@ -1473,41 +1478,51 @@ class BufferedWriter
14731478 " write() argument must be contiguous buffer, not {}" , arg->type ()->name ()));
14741479 }
14751480
1476- // implementation
14771481 if (Buffered<BufferedWriter>::is_closed ()) {
14781482 return Err (value_error (" write to closed file" ));
14791483 }
14801484
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);
1485+ const char *start = static_cast <char *>(buffer.buf ->get_buffer ());
1486+ const auto len = static_cast <size_t >(buffer.len );
1487+
1488+ // fast path: the data fits in the buffer's free space, so it is only buffered and does
1489+ // not reach the raw stream
1490+ if (m_buffered_bytes + len <= m_buffer_size) {
1491+ const auto written = this ->buffer ->sputn (start, len);
1492+ ASSERT (static_cast <size_t >(written) == len);
1493+ m_buffered_bytes += len;
1494+ return PyInteger::create (len);
14861495 }
14871496
1488- // flush our internal buffer to raw
1497+ // the data does not fit: drain our buffer to raw first
14891498 if (auto r = flush (); r.is_err ()) { return r; }
14901499
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
1500+ // write oversized data directly to raw, bypassing the buffer
1501+ size_t written = 0 ;
1502+ size_t remaining = len;
1503+ while (remaining > m_buffer_size) {
15001504 auto r = raw_write (start + written, remaining);
15011505 if (r.is_err ()) { return r; }
15021506 ASSERT (as<PyInteger>(r.unwrap ()));
1503- const auto count = as<PyInteger>(r.unwrap ())->as_size_t ();
1504- written += count;
1505- remaining -= count;
1507+ const auto &count = as<PyInteger>(r.unwrap ())->as_big_int ();
1508+ if (count < 0 || count > remaining) {
1509+ return Err (
1510+ os_error (" raw write() returned invalid length {} (should have been "
1511+ " between 0 and {})" ,
1512+ count.get_str (),
1513+ remaining));
1514+ }
1515+ written += count.get_ui ();
1516+ remaining -= count.get_ui ();
15061517 }
15071518
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 );
1519+ // buffer the tail, which is now guaranteed to fit
1520+ if (remaining > 0 ) {
1521+ const auto count = this ->buffer ->sputn (start + written, remaining);
1522+ ASSERT (static_cast <size_t >(count) == remaining);
1523+ m_buffered_bytes += remaining;
1524+ written += remaining;
1525+ }
15111526
15121527 return PyInteger::create (written);
15131528 }
@@ -1538,6 +1553,10 @@ class BufferedWriter
15381553 PyObject::visit_graph (visitor);
15391554 visit_graph_buffered (visitor);
15401555 }
1556+
1557+ private:
1558+ size_t m_buffer_size{ s_default_buffer_size };
1559+ size_t m_buffered_bytes{ 0 };
15411560};
15421561
15431562template <> PyResult<PyObject *> Buffered<BufferedWriter>::flush_and_rewind()
@@ -1568,6 +1587,7 @@ PyResult<PyObject *> BufferedWriter::flush()
15681587 return r;
15691588 }
15701589 }
1590+ m_buffered_bytes = 0 ;
15711591 return Ok (py_none ());
15721592}
15731593
@@ -3462,9 +3482,8 @@ PyModule *io_module()
34623482 return open (arg0, " rb" );
34633483 }));
34643484
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 });
3485+ s_io_module->add_symbol (PyString::create (" DEFAULT_BUFFER_SIZE" ).unwrap (),
3486+ Number{ static_cast <int64_t >(s_default_buffer_size) });
34683487
34693488 // >> type("UnsupportedOperation", (_io.OSError, ValueError), {})
34703489 auto unsupported_operation_type = types::type ()->call (
0 commit comments