[mypyc] Add big endian read/write functions to librt.strings

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

@@ -21,8 +21,14 @@ class StringWriter:
     def __getitem__(self, i: i64, /) -> i32: ...
 
 def write_i16_le(b: BytesWriter, n: i16, /) -> None: ...
+def write_i16_be(b: BytesWriter, n: i16, /) -> None: ...
 def read_i16_le(b: bytes, index: i64, /) -> i16: ...
+def read_i16_be(b: bytes, index: i64, /) -> i16: ...
 def write_i32_le(b: BytesWriter, n: i32, /) -> None: ...
+def write_i32_be(b: BytesWriter, n: i32, /) -> None: ...
 def read_i32_le(b: bytes, index: i64, /) -> i32: ...
+def read_i32_be(b: bytes, index: i64, /) -> i32: ...
 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: ...

mypyc/lib-rt/byteswriter_extra_ops.h

@@ -74,6 +74,15 @@ CPyBytesWriter_WriteI16LE(PyObject *obj, int16_t value) {
     return CPY_NONE;
 }
 
+static inline char
+CPyBytesWriter_WriteI16BE(PyObject *obj, int16_t value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    if (!CPyBytesWriter_EnsureSize(self, 2))
+        return CPY_NONE_ERROR;
+    BytesWriter_WriteI16BEUnsafe(self, value);
+    return CPY_NONE;
+}
+
 static inline char
 CPyBytesWriter_WriteI32LE(PyObject *obj, int32_t value) {
     BytesWriterObject *self = (BytesWriterObject *)obj;
@@ -83,6 +92,15 @@ CPyBytesWriter_WriteI32LE(PyObject *obj, int32_t value) {
     return CPY_NONE;
 }
 
+static inline char
+CPyBytesWriter_WriteI32BE(PyObject *obj, int32_t value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    if (!CPyBytesWriter_EnsureSize(self, 4))
+        return CPY_NONE_ERROR;
+    BytesWriter_WriteI32BEUnsafe(self, value);
+    return CPY_NONE;
+}
+
 static inline char
 CPyBytesWriter_WriteI64LE(PyObject *obj, int64_t value) {
     BytesWriterObject *self = (BytesWriterObject *)obj;
@@ -92,7 +110,16 @@ CPyBytesWriter_WriteI64LE(PyObject *obj, int64_t value) {
     return CPY_NONE;
 }
 
-// Bytes: Read integer operations (little-endian)
+static inline char
+CPyBytesWriter_WriteI64BE(PyObject *obj, int64_t value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    if (!CPyBytesWriter_EnsureSize(self, 8))
+        return CPY_NONE_ERROR;
+    BytesWriter_WriteI64BEUnsafe(self, value);
+    return CPY_NONE;
+}
+
+// Bytes: Read integer operations
 
 // Helper function for bytes read error handling (negative index or out of range)
 void CPyBytes_ReadError(int64_t index, Py_ssize_t size);
@@ -109,6 +136,30 @@ CPyBytes_ReadI16LE(PyObject *bytes_obj, int64_t index) {
     return CPyBytes_ReadI16LEUnsafe(data + index);
 }
 
+static inline int16_t
+CPyBytes_ReadI16BE(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 - 2)) {
+        CPyBytes_ReadError(index, size);
+        return CPY_LL_INT_ERROR;
+    }
+    const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+    return CPyBytes_ReadI16BEUnsafe(data + index);
+}
+
+static inline int32_t
+CPyBytes_ReadI32BE(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_LL_INT_ERROR;
+    }
+    const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+    return CPyBytes_ReadI32BEUnsafe(data + index);
+}
+
 static inline int32_t
 CPyBytes_ReadI32LE(PyObject *bytes_obj, int64_t index) {
     // bytes_obj type is enforced by mypyc
@@ -133,6 +184,18 @@ CPyBytes_ReadI64LE(PyObject *bytes_obj, int64_t index) {
     return CPyBytes_ReadI64LEUnsafe(data + index);
 }
 
+static inline int64_t
+CPyBytes_ReadI64BE(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_LL_INT_ERROR;
+    }
+    const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+    return CPyBytes_ReadI64BEUnsafe(data + index);
+}
+
 #endif // MYPYC_EXPERIMENTAL
 
 #endif

mypyc/lib-rt/strings/librt_strings.c

@@ -894,6 +894,20 @@ write_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
     Py_RETURN_NONE;
 }
 
+static PyObject*
+write_i16_be(PyObject *module, PyObject *const *args, size_t nargs) {
+    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i16_be");
+    if (bw == NULL)
+        return NULL;
+    int16_t unboxed = CPyLong_AsInt16(args[1]);
+    if (unlikely(unboxed == CPY_LL_INT_ERROR && PyErr_Occurred()))
+        return NULL;
+    if (unlikely(!ensure_bytes_writer_size(bw, 2)))
+        return NULL;
+    BytesWriter_WriteI16BEUnsafe(bw, unboxed);
+    Py_RETURN_NONE;
+}
+
 static PyObject*
 read_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
     int64_t index;
@@ -903,6 +917,15 @@ read_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
     return PyLong_FromLong(CPyBytes_ReadI16LEUnsafe(data + index));
 }
 
+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);
+    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");
@@ -917,6 +940,20 @@ write_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
     Py_RETURN_NONE;
 }
 
+static PyObject*
+write_i32_be(PyObject *module, PyObject *const *args, size_t nargs) {
+    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i32_be");
+    if (bw == NULL)
+        return NULL;
+    int32_t unboxed = CPyLong_AsInt32(args[1]);
+    if (unlikely(unboxed == CPY_LL_INT_ERROR && PyErr_Occurred()))
+        return NULL;
+    if (unlikely(!ensure_bytes_writer_size(bw, 4)))
+        return NULL;
+    BytesWriter_WriteI32BEUnsafe(bw, unboxed);
+    Py_RETURN_NONE;
+}
+
 static PyObject*
 read_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
     int64_t index;
@@ -926,6 +963,15 @@ read_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
     return PyLong_FromLong(CPyBytes_ReadI32LEUnsafe(data + index));
 }
 
+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);
+    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");
@@ -940,6 +986,20 @@ write_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
     Py_RETURN_NONE;
 }
 
+static PyObject*
+write_i64_be(PyObject *module, PyObject *const *args, size_t nargs) {
+    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i64_be");
+    if (bw == NULL)
+        return NULL;
+    int64_t unboxed = CPyLong_AsInt64(args[1]);
+    if (unlikely(unboxed == CPY_LL_INT_ERROR && PyErr_Occurred()))
+        return NULL;
+    if (unlikely(!ensure_bytes_writer_size(bw, 8)))
+        return NULL;
+    BytesWriter_WriteI64BEUnsafe(bw, unboxed);
+    Py_RETURN_NONE;
+}
+
 static PyObject*
 read_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
     int64_t index;
@@ -949,28 +1009,55 @@ read_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
     return PyLong_FromLongLong(CPyBytes_ReadI64LEUnsafe(data + index));
 }
 
+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);
+    if (data == NULL)
+        return NULL;
+    return PyLong_FromLongLong(CPyBytes_ReadI64BEUnsafe(data + index));
+}
+
 #endif
 
 static PyMethodDef librt_strings_module_methods[] = {
 #ifdef MYPYC_EXPERIMENTAL
     {"write_i16_le", (PyCFunction) write_i16_le, METH_FASTCALL,
      PyDoc_STR("Write a 16-bit signed integer to BytesWriter in little-endian format")
     },
+    {"write_i16_be", (PyCFunction) write_i16_be, METH_FASTCALL,
+     PyDoc_STR("Write a 16-bit signed integer to BytesWriter in big-endian format")
+    },
     {"read_i16_le", (PyCFunction) read_i16_le, METH_FASTCALL,
      PyDoc_STR("Read a 16-bit signed integer from bytes in little-endian format")
     },
+    {"read_i16_be", (PyCFunction) read_i16_be, METH_FASTCALL,
+     PyDoc_STR("Read a 16-bit signed integer from bytes in big-endian format")
+    },
     {"write_i32_le", (PyCFunction) write_i32_le, METH_FASTCALL,
      PyDoc_STR("Write a 32-bit signed integer to BytesWriter in little-endian format")
     },
+    {"write_i32_be", (PyCFunction) write_i32_be, METH_FASTCALL,
+     PyDoc_STR("Write a 32-bit signed integer to BytesWriter in big-endian format")
+    },
     {"read_i32_le", (PyCFunction) read_i32_le, METH_FASTCALL,
      PyDoc_STR("Read a 32-bit signed integer from bytes in little-endian format")
     },
+    {"read_i32_be", (PyCFunction) read_i32_be, METH_FASTCALL,
+     PyDoc_STR("Read a 32-bit signed integer from bytes in big-endian format")
+    },
     {"write_i64_le", (PyCFunction) write_i64_le, METH_FASTCALL,
      PyDoc_STR("Write a 64-bit signed integer to BytesWriter in little-endian format")
     },
+    {"write_i64_be", (PyCFunction) write_i64_be, METH_FASTCALL,
+     PyDoc_STR("Write a 64-bit signed integer to BytesWriter in big-endian format")
+    },
     {"read_i64_le", (PyCFunction) read_i64_le, METH_FASTCALL,
      PyDoc_STR("Read a 64-bit signed integer from bytes in little-endian format")
     },
+    {"read_i64_be", (PyCFunction) read_i64_be, METH_FASTCALL,
+     PyDoc_STR("Read a 64-bit signed integer from bytes in big-endian format")
+    },
 #endif
     {NULL, NULL, 0, NULL}
 };