[mypyc] Add float read/write functions to librt.strings (#20778)

Support both 32-bit and 64-bit float formats, and also support little
endian and big endian (e.g. `write_f64_le`).

I think the read/write function feature set now covers all the essential
functionality. We can consider adding `u8` variants in the future,
though they can already be supported using just `append` and get item.

Follow-up to #20772.

I used coding agent assist significantly.

Author: JukkaL

mypy/typeshed/stubs/librt/librt/strings.pyi

@@ -32,3 +32,11 @@ def write_i64_le(b: BytesWriter, n: i64, /) -> None: ...
 def write_i64_be(b: BytesWriter, n: i64, /) -> None: ...
 def read_i64_le(b: bytes, index: i64, /) -> i64: ...
 def read_i64_be(b: bytes, index: i64, /) -> i64: ...
+def write_f32_le(b: BytesWriter, n: float, /) -> None: ...
+def write_f32_be(b: BytesWriter, n: float, /) -> None: ...
+def read_f32_le(b: bytes, index: i64, /) -> float: ...
+def read_f32_be(b: bytes, index: i64, /) -> float: ...
+def write_f64_le(b: BytesWriter, n: float, /) -> None: ...
+def write_f64_be(b: BytesWriter, n: float, /) -> None: ...
+def read_f64_le(b: bytes, index: i64, /) -> float: ...
+def read_f64_be(b: bytes, index: i64, /) -> float: ...

mypyc/lib-rt/byteswriter_extra_ops.h

@@ -119,6 +119,44 @@ CPyBytesWriter_WriteI64BE(PyObject *obj, int64_t value) {
     return CPY_NONE;
 }
 
+// BytesWriter: Write float operations
+
+static inline char
+CPyBytesWriter_WriteF32LE(PyObject *obj, double value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    if (!CPyBytesWriter_EnsureSize(self, 4))
+        return CPY_NONE_ERROR;
+    BytesWriter_WriteF32LEUnsafe(self, (float)value);
+    return CPY_NONE;
+}
+
+static inline char
+CPyBytesWriter_WriteF32BE(PyObject *obj, double value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    if (!CPyBytesWriter_EnsureSize(self, 4))
+        return CPY_NONE_ERROR;
+    BytesWriter_WriteF32BEUnsafe(self, (float)value);
+    return CPY_NONE;
+}
+
+static inline char
+CPyBytesWriter_WriteF64LE(PyObject *obj, double value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    if (!CPyBytesWriter_EnsureSize(self, 8))
+        return CPY_NONE_ERROR;
+    BytesWriter_WriteF64LEUnsafe(self, value);
+    return CPY_NONE;
+}
+
+static inline char
+CPyBytesWriter_WriteF64BE(PyObject *obj, double value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    if (!CPyBytesWriter_EnsureSize(self, 8))
+        return CPY_NONE_ERROR;
+    BytesWriter_WriteF64BEUnsafe(self, value);
+    return CPY_NONE;
+}
+
 // Bytes: Read integer operations
 
 // Helper function for bytes read error handling (negative index or out of range)
@@ -196,6 +234,56 @@ CPyBytes_ReadI64BE(PyObject *bytes_obj, int64_t index) {
     return CPyBytes_ReadI64BEUnsafe(data + index);
 }
 
+// Bytes: Read float operations
+
+static inline double
+CPyBytes_ReadF32LE(PyObject *bytes_obj, int64_t index) {
+    // bytes_obj type is enforced by mypyc
+    Py_ssize_t size = PyBytes_GET_SIZE(bytes_obj);
+    if (unlikely(index < 0 || index > size - 4)) {
+        CPyBytes_ReadError(index, size);
+        return CPY_FLOAT_ERROR;
+    }
+    const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+    return (double)CPyBytes_ReadF32LEUnsafe(data + index);
+}
+
+static inline double
+CPyBytes_ReadF32BE(PyObject *bytes_obj, int64_t index) {
+    // bytes_obj type is enforced by mypyc
+    Py_ssize_t size = PyBytes_GET_SIZE(bytes_obj);
+    if (unlikely(index < 0 || index > size - 4)) {
+        CPyBytes_ReadError(index, size);
+        return CPY_FLOAT_ERROR;
+    }
+    const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+    return (double)CPyBytes_ReadF32BEUnsafe(data + index);
+}
+
+static inline double
+CPyBytes_ReadF64LE(PyObject *bytes_obj, int64_t index) {
+    // bytes_obj type is enforced by mypyc
+    Py_ssize_t size = PyBytes_GET_SIZE(bytes_obj);
+    if (unlikely(index < 0 || index > size - 8)) {
+        CPyBytes_ReadError(index, size);
+        return CPY_FLOAT_ERROR;
+    }
+    const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+    return CPyBytes_ReadF64LEUnsafe(data + index);
+}
+
+static inline double
+CPyBytes_ReadF64BE(PyObject *bytes_obj, int64_t index) {
+    // bytes_obj type is enforced by mypyc
+    Py_ssize_t size = PyBytes_GET_SIZE(bytes_obj);
+    if (unlikely(index < 0 || index > size - 8)) {
+        CPyBytes_ReadError(index, size);
+        return CPY_FLOAT_ERROR;
+    }
+    const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+    return CPyBytes_ReadF64BEUnsafe(data + index);
+}
+
 #endif // MYPYC_EXPERIMENTAL
 
 #endif

mypyc/lib-rt/strings/librt_strings.c

@@ -831,9 +831,9 @@ StringWriter_len_internal(PyObject *self) {
 
 // End of StringWriter
 
-// Helper for write_i*_le/be functions - validates args and returns BytesWriter
+// Helper for write_*_le/be functions - validates args and returns BytesWriter
 static inline BytesWriterObject *
-parse_write_int_args(PyObject *const *args, size_t nargs, const char *func_name) {
+parse_write_args(PyObject *const *args, size_t nargs, const char *func_name) {
     if (unlikely(nargs != 2)) {
         PyErr_Format(PyExc_TypeError,
                      "%s() takes exactly 2 arguments (%zu given)", func_name, nargs);
@@ -846,10 +846,10 @@ parse_write_int_args(PyObject *const *args, size_t nargs, const char *func_name)
     return (BytesWriterObject *)writer;
 }
 
-// Helper for read_i*_le/be functions - validates args and returns data pointer
+// Helper for read_*_le/be functions - validates args and returns data pointer
 // Returns NULL on error, sets *out_index to the validated index on success
 static inline const unsigned char *
-parse_read_int_args(PyObject *const *args, size_t nargs, const char *func_name,
+parse_read_args(PyObject *const *args, size_t nargs, const char *func_name,
                     Py_ssize_t num_bytes, int64_t *out_index) {
     if (unlikely(nargs != 2)) {
         PyErr_Format(PyExc_TypeError,
@@ -882,7 +882,7 @@ parse_read_int_args(PyObject *const *args, size_t nargs, const char *func_name,
 
 static PyObject*
 write_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
-    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i16_le");
+    BytesWriterObject *bw = parse_write_args(args, nargs, "write_i16_le");
     if (bw == NULL)
         return NULL;
     int16_t unboxed = CPyLong_AsInt16(args[1]);
@@ -896,7 +896,7 @@ write_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
 
 static PyObject*
 write_i16_be(PyObject *module, PyObject *const *args, size_t nargs) {
-    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i16_be");
+    BytesWriterObject *bw = parse_write_args(args, nargs, "write_i16_be");
     if (bw == NULL)
         return NULL;
     int16_t unboxed = CPyLong_AsInt16(args[1]);
@@ -911,7 +911,7 @@ write_i16_be(PyObject *module, PyObject *const *args, size_t nargs) {
 static PyObject*
 read_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
     int64_t index;
-    const unsigned char *data = parse_read_int_args(args, nargs, "read_i16_le", 2, &index);
+    const unsigned char *data = parse_read_args(args, nargs, "read_i16_le", 2, &index);
     if (data == NULL)
         return NULL;
     return PyLong_FromLong(CPyBytes_ReadI16LEUnsafe(data + index));
@@ -920,15 +920,15 @@ read_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
 static PyObject*
 read_i16_be(PyObject *module, PyObject *const *args, size_t nargs) {
     int64_t index;
-    const unsigned char *data = parse_read_int_args(args, nargs, "read_i16_be", 2, &index);
+    const unsigned char *data = parse_read_args(args, nargs, "read_i16_be", 2, &index);
     if (data == NULL)
         return NULL;
     return PyLong_FromLong(CPyBytes_ReadI16BEUnsafe(data + index));
 }
 
 static PyObject*
 write_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
-    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i32_le");
+    BytesWriterObject *bw = parse_write_args(args, nargs, "write_i32_le");
     if (bw == NULL)
         return NULL;
     int32_t unboxed = CPyLong_AsInt32(args[1]);
@@ -942,7 +942,7 @@ write_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
 
 static PyObject*
 write_i32_be(PyObject *module, PyObject *const *args, size_t nargs) {
-    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i32_be");
+    BytesWriterObject *bw = parse_write_args(args, nargs, "write_i32_be");
     if (bw == NULL)
         return NULL;
     int32_t unboxed = CPyLong_AsInt32(args[1]);
@@ -957,7 +957,7 @@ write_i32_be(PyObject *module, PyObject *const *args, size_t nargs) {
 static PyObject*
 read_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
     int64_t index;
-    const unsigned char *data = parse_read_int_args(args, nargs, "read_i32_le", 4, &index);
+    const unsigned char *data = parse_read_args(args, nargs, "read_i32_le", 4, &index);
     if (data == NULL)
         return NULL;
     return PyLong_FromLong(CPyBytes_ReadI32LEUnsafe(data + index));
@@ -966,15 +966,15 @@ read_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
 static PyObject*
 read_i32_be(PyObject *module, PyObject *const *args, size_t nargs) {
     int64_t index;
-    const unsigned char *data = parse_read_int_args(args, nargs, "read_i32_be", 4, &index);
+    const unsigned char *data = parse_read_args(args, nargs, "read_i32_be", 4, &index);
     if (data == NULL)
         return NULL;
     return PyLong_FromLong(CPyBytes_ReadI32BEUnsafe(data + index));
 }
 
 static PyObject*
 write_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
-    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i64_le");
+    BytesWriterObject *bw = parse_write_args(args, nargs, "write_i64_le");
     if (bw == NULL)
         return NULL;
     int64_t unboxed = CPyLong_AsInt64(args[1]);
@@ -988,7 +988,7 @@ write_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
 
 static PyObject*
 write_i64_be(PyObject *module, PyObject *const *args, size_t nargs) {
-    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i64_be");
+    BytesWriterObject *bw = parse_write_args(args, nargs, "write_i64_be");
     if (bw == NULL)
         return NULL;
     int64_t unboxed = CPyLong_AsInt64(args[1]);
@@ -1003,7 +1003,7 @@ write_i64_be(PyObject *module, PyObject *const *args, size_t nargs) {
 static PyObject*
 read_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
     int64_t index;
-    const unsigned char *data = parse_read_int_args(args, nargs, "read_i64_le", 8, &index);
+    const unsigned char *data = parse_read_args(args, nargs, "read_i64_le", 8, &index);
     if (data == NULL)
         return NULL;
     return PyLong_FromLongLong(CPyBytes_ReadI64LEUnsafe(data + index));
@@ -1012,12 +1012,104 @@ read_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
 static PyObject*
 read_i64_be(PyObject *module, PyObject *const *args, size_t nargs) {
     int64_t index;
-    const unsigned char *data = parse_read_int_args(args, nargs, "read_i64_be", 8, &index);
+    const unsigned char *data = parse_read_args(args, nargs, "read_i64_be", 8, &index);
     if (data == NULL)
         return NULL;
     return PyLong_FromLongLong(CPyBytes_ReadI64BEUnsafe(data + index));
 }
 
+static PyObject*
+write_f32_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    BytesWriterObject *bw = parse_write_args(args, nargs, "write_f32_le");
+    if (bw == NULL)
+        return NULL;
+    double unboxed = PyFloat_AsDouble(args[1]);
+    if (unlikely(unboxed == -1.0 && PyErr_Occurred()))
+        return NULL;
+    if (unlikely(!ensure_bytes_writer_size(bw, 4)))
+        return NULL;
+    BytesWriter_WriteF32LEUnsafe(bw, (float)unboxed);
+    Py_RETURN_NONE;
+}
+
+static PyObject*
+write_f32_be(PyObject *module, PyObject *const *args, size_t nargs) {
+    BytesWriterObject *bw = parse_write_args(args, nargs, "write_f32_be");
+    if (bw == NULL)
+        return NULL;
+    double unboxed = PyFloat_AsDouble(args[1]);
+    if (unlikely(unboxed == -1.0 && PyErr_Occurred()))
+        return NULL;
+    if (unlikely(!ensure_bytes_writer_size(bw, 4)))
+        return NULL;
+    BytesWriter_WriteF32BEUnsafe(bw, (float)unboxed);
+    Py_RETURN_NONE;
+}
+
+static PyObject*
+read_f32_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    int64_t index;
+    const unsigned char *data = parse_read_args(args, nargs, "read_f32_le", 4, &index);
+    if (data == NULL)
+        return NULL;
+    return PyFloat_FromDouble((double)CPyBytes_ReadF32LEUnsafe(data + index));
+}
+
+static PyObject*
+read_f32_be(PyObject *module, PyObject *const *args, size_t nargs) {
+    int64_t index;
+    const unsigned char *data = parse_read_args(args, nargs, "read_f32_be", 4, &index);
+    if (data == NULL)
+        return NULL;
+    return PyFloat_FromDouble((double)CPyBytes_ReadF32BEUnsafe(data + index));
+}
+
+static PyObject*
+write_f64_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    BytesWriterObject *bw = parse_write_args(args, nargs, "write_f64_le");
+    if (bw == NULL)
+        return NULL;
+    double unboxed = PyFloat_AsDouble(args[1]);
+    if (unlikely(unboxed == -1.0 && PyErr_Occurred()))
+        return NULL;
+    if (unlikely(!ensure_bytes_writer_size(bw, 8)))
+        return NULL;
+    BytesWriter_WriteF64LEUnsafe(bw, unboxed);
+    Py_RETURN_NONE;
+}
+
+static PyObject*
+write_f64_be(PyObject *module, PyObject *const *args, size_t nargs) {
+    BytesWriterObject *bw = parse_write_args(args, nargs, "write_f64_be");
+    if (bw == NULL)
+        return NULL;
+    double unboxed = PyFloat_AsDouble(args[1]);
+    if (unlikely(unboxed == -1.0 && PyErr_Occurred()))
+        return NULL;
+    if (unlikely(!ensure_bytes_writer_size(bw, 8)))
+        return NULL;
+    BytesWriter_WriteF64BEUnsafe(bw, unboxed);
+    Py_RETURN_NONE;
+}
+
+static PyObject*
+read_f64_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    int64_t index;
+    const unsigned char *data = parse_read_args(args, nargs, "read_f64_le", 8, &index);
+    if (data == NULL)
+        return NULL;
+    return PyFloat_FromDouble(CPyBytes_ReadF64LEUnsafe(data + index));
+}
+
+static PyObject*
+read_f64_be(PyObject *module, PyObject *const *args, size_t nargs) {
+    int64_t index;
+    const unsigned char *data = parse_read_args(args, nargs, "read_f64_be", 8, &index);
+    if (data == NULL)
+        return NULL;
+    return PyFloat_FromDouble(CPyBytes_ReadF64BEUnsafe(data + index));
+}
+
 #endif
 
 static PyMethodDef librt_strings_module_methods[] = {
@@ -1058,6 +1150,30 @@ static PyMethodDef librt_strings_module_methods[] = {
     {"read_i64_be", (PyCFunction) read_i64_be, METH_FASTCALL,
      PyDoc_STR("Read a 64-bit signed integer from bytes in big-endian format")
     },
+    {"write_f32_le", (PyCFunction) write_f32_le, METH_FASTCALL,
+     PyDoc_STR("Write a 32-bit float to BytesWriter in little-endian format")
+    },
+    {"write_f32_be", (PyCFunction) write_f32_be, METH_FASTCALL,
+     PyDoc_STR("Write a 32-bit float to BytesWriter in big-endian format")
+    },
+    {"read_f32_le", (PyCFunction) read_f32_le, METH_FASTCALL,
+     PyDoc_STR("Read a 32-bit float from bytes in little-endian format")
+    },
+    {"read_f32_be", (PyCFunction) read_f32_be, METH_FASTCALL,
+     PyDoc_STR("Read a 32-bit float from bytes in big-endian format")
+    },
+    {"write_f64_le", (PyCFunction) write_f64_le, METH_FASTCALL,
+     PyDoc_STR("Write a 64-bit float to BytesWriter in little-endian format")
+    },
+    {"write_f64_be", (PyCFunction) write_f64_be, METH_FASTCALL,
+     PyDoc_STR("Write a 64-bit float to BytesWriter in big-endian format")
+    },
+    {"read_f64_le", (PyCFunction) read_f64_le, METH_FASTCALL,
+     PyDoc_STR("Read a 64-bit float from bytes in little-endian format")
+    },
+    {"read_f64_be", (PyCFunction) read_f64_be, METH_FASTCALL,
+     PyDoc_STR("Read a 64-bit float from bytes in big-endian format")
+    },
 #endif
     {NULL, NULL, 0, NULL}
 };