From e453d9d9f14ec5eabd1077d32ed8f293faeefb13 Mon Sep 17 00:00:00 2001
From: robert-hh <robert@hammelrath.com>
Date: Mon, 8 Sep 2025 17:09:12 +0200
Subject: [PATCH] py/stream: Support additional arguments for
 mp_stream_write1_obj.

As suggested by @dpgeorge adding a `stream_write_generic()` function.

Tested the different write timeouts with a RP2 using flow control.  Tested
support of the additional arguments of `uart.write()`.

Signed-off-by: robert-hh <robert@hammelrath.com>
---
 py/stream.c | 16 +++++++++-------
 py/stream.h |  2 +-
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/py/stream.c b/py/stream.c
index e41fbdd5da..008c968452 100644
--- a/py/stream.c
+++ b/py/stream.c
@@ -261,7 +261,7 @@ void mp_stream_write_adaptor(void *self, const char *buf, size_t len) {
     mp_stream_write(MP_OBJ_FROM_PTR(self), buf, len, MP_STREAM_RW_WRITE);
 }
 
-static mp_obj_t stream_write_method(size_t n_args, const mp_obj_t *args) {
+static mp_obj_t stream_write_generic(size_t n_args, const mp_obj_t *args, byte flags) {
     mp_buffer_info_t bufinfo;
     mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_READ);
     size_t max_len = (size_t)-1;
@@ -276,16 +276,18 @@ static mp_obj_t stream_write_method(size_t n_args, const mp_obj_t *args) {
         }
     }
     bufinfo.len -= off;
-    return mp_stream_write(args[0], (byte *)bufinfo.buf + off, MIN(bufinfo.len, max_len), MP_STREAM_RW_WRITE);
+    return mp_stream_write(args[0], (byte *)bufinfo.buf + off, MIN(bufinfo.len, max_len), flags);
+}
+
+static mp_obj_t stream_write_method(size_t n_args, const mp_obj_t *args) {
+    return stream_write_generic(n_args, args, MP_STREAM_RW_WRITE);
 }
 MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_write_obj, 2, 4, stream_write_method);
 
-static mp_obj_t stream_write1_method(mp_obj_t self_in, mp_obj_t arg) {
-    mp_buffer_info_t bufinfo;
-    mp_get_buffer_raise(arg, &bufinfo, MP_BUFFER_READ);
-    return mp_stream_write(self_in, bufinfo.buf, bufinfo.len, MP_STREAM_RW_WRITE | MP_STREAM_RW_ONCE);
+static mp_obj_t stream_write1_method(size_t n_args, const mp_obj_t *args) {
+    return stream_write_generic(n_args, args, MP_STREAM_RW_WRITE | MP_STREAM_RW_ONCE);
 }
-MP_DEFINE_CONST_FUN_OBJ_2(mp_stream_write1_obj, stream_write1_method);
+MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_write1_obj, 2, 4, stream_write1_method);
 
 static mp_obj_t stream_readinto_generic(size_t n_args, const mp_obj_t *args, byte flags) {
     mp_buffer_info_t bufinfo;
diff --git a/py/stream.h b/py/stream.h
index c3606e8ece..5a3462d5a3 100644
--- a/py/stream.h
+++ b/py/stream.h
@@ -83,7 +83,7 @@ MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_readinto1_obj);
 MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_unbuffered_readline_obj);
 MP_DECLARE_CONST_FUN_OBJ_1(mp_stream_unbuffered_readlines_obj);
 MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_write_obj);
-MP_DECLARE_CONST_FUN_OBJ_2(mp_stream_write1_obj);
+MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_write1_obj);
 MP_DECLARE_CONST_FUN_OBJ_1(mp_stream_close_obj);
 MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream___exit___obj);
 MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_seek_obj);