Skip to content

Commit 2955e3c

Browse files
JukkaLJukkaL
authored
[mypyc] Add big endian read/write functions to librt.strings (#20772)
This is pretty straightforward -- add big endian variants of existing functions. Follow-up to #20757. I'll add float read/write functions in a follow-up PR. Also restructure primitive definitions to have more consistent order. I used coding agent assist to generate the repetitive parts and to refactor the code and tests.
1 parent 3143ad5 commit 2955e3c

6 files changed

Lines changed: 660 additions & 203 deletions

File tree

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,14 @@ class StringWriter:
2121
def __getitem__(self, i: i64, /) -> i32: ...
2222

2323
def write_i16_le(b: BytesWriter, n: i16, /) -> None: ...
24+
def write_i16_be(b: BytesWriter, n: i16, /) -> None: ...
2425
def read_i16_le(b: bytes, index: i64, /) -> i16: ...
26+
def read_i16_be(b: bytes, index: i64, /) -> i16: ...
2527
def write_i32_le(b: BytesWriter, n: i32, /) -> None: ...
28+
def write_i32_be(b: BytesWriter, n: i32, /) -> None: ...
2629
def read_i32_le(b: bytes, index: i64, /) -> i32: ...
30+
def read_i32_be(b: bytes, index: i64, /) -> i32: ...
2731
def write_i64_le(b: BytesWriter, n: i64, /) -> None: ...
32+
def write_i64_be(b: BytesWriter, n: i64, /) -> None: ...
2833
def read_i64_le(b: bytes, index: i64, /) -> i64: ...
34+
def read_i64_be(b: bytes, index: i64, /) -> i64: ...

‎mypyc/lib-rt/byteswriter_extra_ops.h‎

Lines changed: 64 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,15 @@ CPyBytesWriter_WriteI16LE(PyObject *obj, int16_t value) {
7474
return CPY_NONE;
7575
}
7676

77+
static inline char
78+
CPyBytesWriter_WriteI16BE(PyObject *obj, int16_t value) {
79+
BytesWriterObject *self = (BytesWriterObject *)obj;
80+
if (!CPyBytesWriter_EnsureSize(self, 2))
81+
return CPY_NONE_ERROR;
82+
BytesWriter_WriteI16BEUnsafe(self, value);
83+
return CPY_NONE;
84+
}
85+
7786
static inline char
7887
CPyBytesWriter_WriteI32LE(PyObject *obj, int32_t value) {
7988
BytesWriterObject *self = (BytesWriterObject *)obj;
@@ -83,6 +92,15 @@ CPyBytesWriter_WriteI32LE(PyObject *obj, int32_t value) {
8392
return CPY_NONE;
8493
}
8594

95+
static inline char
96+
CPyBytesWriter_WriteI32BE(PyObject *obj, int32_t value) {
97+
BytesWriterObject *self = (BytesWriterObject *)obj;
98+
if (!CPyBytesWriter_EnsureSize(self, 4))
99+
return CPY_NONE_ERROR;
100+
BytesWriter_WriteI32BEUnsafe(self, value);
101+
return CPY_NONE;
102+
}
103+
86104
static inline char
87105
CPyBytesWriter_WriteI64LE(PyObject *obj, int64_t value) {
88106
BytesWriterObject *self = (BytesWriterObject *)obj;
@@ -92,7 +110,16 @@ CPyBytesWriter_WriteI64LE(PyObject *obj, int64_t value) {
92110
return CPY_NONE;
93111
}
94112

95-
// Bytes: Read integer operations (little-endian)
113+
static inline char
114+
CPyBytesWriter_WriteI64BE(PyObject *obj, int64_t value) {
115+
BytesWriterObject *self = (BytesWriterObject *)obj;
116+
if (!CPyBytesWriter_EnsureSize(self, 8))
117+
return CPY_NONE_ERROR;
118+
BytesWriter_WriteI64BEUnsafe(self, value);
119+
return CPY_NONE;
120+
}
121+
122+
// Bytes: Read integer operations
96123

97124
// Helper function for bytes read error handling (negative index or out of range)
98125
void CPyBytes_ReadError(int64_t index, Py_ssize_t size);
@@ -109,6 +136,30 @@ CPyBytes_ReadI16LE(PyObject *bytes_obj, int64_t index) {
109136
return CPyBytes_ReadI16LEUnsafe(data + index);
110137
}
111138

139+
static inline int16_t
140+
CPyBytes_ReadI16BE(PyObject *bytes_obj, int64_t index) {
141+
// bytes_obj type is enforced by mypyc
142+
Py_ssize_t size = PyBytes_GET_SIZE(bytes_obj);
143+
if (unlikely(index < 0 || index > size - 2)) {
144+
CPyBytes_ReadError(index, size);
145+
return CPY_LL_INT_ERROR;
146+
}
147+
const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
148+
return CPyBytes_ReadI16BEUnsafe(data + index);
149+
}
150+
151+
static inline int32_t
152+
CPyBytes_ReadI32BE(PyObject *bytes_obj, int64_t index) {
153+
// bytes_obj type is enforced by mypyc
154+
Py_ssize_t size = PyBytes_GET_SIZE(bytes_obj);
155+
if (unlikely(index < 0 || index > size - 4)) {
156+
CPyBytes_ReadError(index, size);
157+
return CPY_LL_INT_ERROR;
158+
}
159+
const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
160+
return CPyBytes_ReadI32BEUnsafe(data + index);
161+
}
162+
112163
static inline int32_t
113164
CPyBytes_ReadI32LE(PyObject *bytes_obj, int64_t index) {
114165
// bytes_obj type is enforced by mypyc
@@ -133,6 +184,18 @@ CPyBytes_ReadI64LE(PyObject *bytes_obj, int64_t index) {
133184
return CPyBytes_ReadI64LEUnsafe(data + index);
134185
}
135186

187+
static inline int64_t
188+
CPyBytes_ReadI64BE(PyObject *bytes_obj, int64_t index) {
189+
// bytes_obj type is enforced by mypyc
190+
Py_ssize_t size = PyBytes_GET_SIZE(bytes_obj);
191+
if (unlikely(index < 0 || index > size - 8)) {
192+
CPyBytes_ReadError(index, size);
193+
return CPY_LL_INT_ERROR;
194+
}
195+
const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
196+
return CPyBytes_ReadI64BEUnsafe(data + index);
197+
}
198+
136199
#endif // MYPYC_EXPERIMENTAL
137200

138201
#endif

‎mypyc/lib-rt/strings/librt_strings.c‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -894,6 +894,20 @@ write_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
894894
Py_RETURN_NONE;
895895
}
896896

897+
static PyObject*
898+
write_i16_be(PyObject *module, PyObject *const *args, size_t nargs) {
899+
BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i16_be");
900+
if (bw == NULL)
901+
return NULL;
902+
int16_t unboxed = CPyLong_AsInt16(args[1]);
903+
if (unlikely(unboxed == CPY_LL_INT_ERROR && PyErr_Occurred()))
904+
return NULL;
905+
if (unlikely(!ensure_bytes_writer_size(bw, 2)))
906+
return NULL;
907+
BytesWriter_WriteI16BEUnsafe(bw, unboxed);
908+
Py_RETURN_NONE;
909+
}
910+
897911
static PyObject*
898912
read_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
899913
int64_t index;
@@ -903,6 +917,15 @@ read_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
903917
return PyLong_FromLong(CPyBytes_ReadI16LEUnsafe(data + index));
904918
}
905919

920+
static PyObject*
921+
read_i16_be(PyObject *module, PyObject *const *args, size_t nargs) {
922+
int64_t index;
923+
const unsigned char *data = parse_read_int_args(args, nargs, "read_i16_be", 2, &index);
924+
if (data == NULL)
925+
return NULL;
926+
return PyLong_FromLong(CPyBytes_ReadI16BEUnsafe(data + index));
927+
}
928+
906929
static PyObject*
907930
write_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
908931
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) {
917940
Py_RETURN_NONE;
918941
}
919942

943+
static PyObject*
944+
write_i32_be(PyObject *module, PyObject *const *args, size_t nargs) {
945+
BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i32_be");
946+
if (bw == NULL)
947+
return NULL;
948+
int32_t unboxed = CPyLong_AsInt32(args[1]);
949+
if (unlikely(unboxed == CPY_LL_INT_ERROR && PyErr_Occurred()))
950+
return NULL;
951+
if (unlikely(!ensure_bytes_writer_size(bw, 4)))
952+
return NULL;
953+
BytesWriter_WriteI32BEUnsafe(bw, unboxed);
954+
Py_RETURN_NONE;
955+
}
956+
920957
static PyObject*
921958
read_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
922959
int64_t index;
@@ -926,6 +963,15 @@ read_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
926963
return PyLong_FromLong(CPyBytes_ReadI32LEUnsafe(data + index));
927964
}
928965

966+
static PyObject*
967+
read_i32_be(PyObject *module, PyObject *const *args, size_t nargs) {
968+
int64_t index;
969+
const unsigned char *data = parse_read_int_args(args, nargs, "read_i32_be", 4, &index);
970+
if (data == NULL)
971+
return NULL;
972+
return PyLong_FromLong(CPyBytes_ReadI32BEUnsafe(data + index));
973+
}
974+
929975
static PyObject*
930976
write_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
931977
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) {
940986
Py_RETURN_NONE;
941987
}
942988

989+
static PyObject*
990+
write_i64_be(PyObject *module, PyObject *const *args, size_t nargs) {
991+
BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i64_be");
992+
if (bw == NULL)
993+
return NULL;
994+
int64_t unboxed = CPyLong_AsInt64(args[1]);
995+
if (unlikely(unboxed == CPY_LL_INT_ERROR && PyErr_Occurred()))
996+
return NULL;
997+
if (unlikely(!ensure_bytes_writer_size(bw, 8)))
998+
return NULL;
999+
BytesWriter_WriteI64BEUnsafe(bw, unboxed);
1000+
Py_RETURN_NONE;
1001+
}
1002+
9431003
static PyObject*
9441004
read_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
9451005
int64_t index;
@@ -949,28 +1009,55 @@ read_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
9491009
return PyLong_FromLongLong(CPyBytes_ReadI64LEUnsafe(data + index));
9501010
}
9511011

1012+
static PyObject*
1013+
read_i64_be(PyObject *module, PyObject *const *args, size_t nargs) {
1014+
int64_t index;
1015+
const unsigned char *data = parse_read_int_args(args, nargs, "read_i64_be", 8, &index);
1016+
if (data == NULL)
1017+
return NULL;
1018+
return PyLong_FromLongLong(CPyBytes_ReadI64BEUnsafe(data + index));
1019+
}
1020+
9521021
#endif
9531022

9541023
static PyMethodDef librt_strings_module_methods[] = {
9551024
#ifdef MYPYC_EXPERIMENTAL
9561025
{"write_i16_le", (PyCFunction) write_i16_le, METH_FASTCALL,
9571026
PyDoc_STR("Write a 16-bit signed integer to BytesWriter in little-endian format")
9581027
},
1028+
{"write_i16_be", (PyCFunction) write_i16_be, METH_FASTCALL,
1029+
PyDoc_STR("Write a 16-bit signed integer to BytesWriter in big-endian format")
1030+
},
9591031
{"read_i16_le", (PyCFunction) read_i16_le, METH_FASTCALL,
9601032
PyDoc_STR("Read a 16-bit signed integer from bytes in little-endian format")
9611033
},
1034+
{"read_i16_be", (PyCFunction) read_i16_be, METH_FASTCALL,
1035+
PyDoc_STR("Read a 16-bit signed integer from bytes in big-endian format")
1036+
},
9621037
{"write_i32_le", (PyCFunction) write_i32_le, METH_FASTCALL,
9631038
PyDoc_STR("Write a 32-bit signed integer to BytesWriter in little-endian format")
9641039
},
1040+
{"write_i32_be", (PyCFunction) write_i32_be, METH_FASTCALL,
1041+
PyDoc_STR("Write a 32-bit signed integer to BytesWriter in big-endian format")
1042+
},
9651043
{"read_i32_le", (PyCFunction) read_i32_le, METH_FASTCALL,
9661044
PyDoc_STR("Read a 32-bit signed integer from bytes in little-endian format")
9671045
},
1046+
{"read_i32_be", (PyCFunction) read_i32_be, METH_FASTCALL,
1047+
PyDoc_STR("Read a 32-bit signed integer from bytes in big-endian format")
1048+
},
9681049
{"write_i64_le", (PyCFunction) write_i64_le, METH_FASTCALL,
9691050
PyDoc_STR("Write a 64-bit signed integer to BytesWriter in little-endian format")
9701051
},
1052+
{"write_i64_be", (PyCFunction) write_i64_be, METH_FASTCALL,
1053+
PyDoc_STR("Write a 64-bit signed integer to BytesWriter in big-endian format")
1054+
},
9711055
{"read_i64_le", (PyCFunction) read_i64_le, METH_FASTCALL,
9721056
PyDoc_STR("Read a 64-bit signed integer from bytes in little-endian format")
9731057
},
1058+
{"read_i64_be", (PyCFunction) read_i64_be, METH_FASTCALL,
1059+
PyDoc_STR("Read a 64-bit signed integer from bytes in big-endian format")
1060+
},
9741061
#endif
9751062
{NULL, NULL, 0, NULL}
9761063
};

0 commit comments

Comments
 (0)