Skip to content

Commit 32f7c62

Browse files
vstinnervstinner
committed
PEP 782: Add PyBytesWriter C API
1 parent 632d1aa commit 32f7c62

File tree

2 files changed

+348
-0
lines changed

2 files changed

+348
-0
lines changed

‎pythoncapi_compat.h

Lines changed: 250 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2199,6 +2199,256 @@ PyConfig_GetInt(const char *name, int *value)
21992199
#endif // PY_VERSION_HEX > 0x03090000 && !defined(PYPY_VERSION)
22002200

22012201

2202+
#if PY_VERSION_HEX < 0x030E00A7
2203+
typedef struct PyBytesWriter {
2204+
char small_buffer[256];
2205+
PyObject *obj;
2206+
Py_ssize_t size;
2207+
} PyBytesWriter;
2208+
2209+
static inline Py_ssize_t
2210+
_PyBytesWriter_GetAllocated(PyBytesWriter *writer)
2211+
{
2212+
if (writer->obj == NULL) {
2213+
return sizeof(writer->small_buffer);
2214+
}
2215+
else {
2216+
return PyBytes_GET_SIZE(writer->obj);
2217+
}
2218+
}
2219+
2220+
2221+
static inline int
2222+
_PyBytesWriter_Resize_impl(PyBytesWriter *writer, Py_ssize_t size,
2223+
int overallocate)
2224+
{
2225+
assert(size >= 0);
2226+
2227+
if (size <= _PyBytesWriter_GetAllocated(writer)) {
2228+
return 0;
2229+
}
2230+
2231+
if (overallocate) {
2232+
#ifdef MS_WINDOWS
2233+
/* On Windows, overallocate by 50% is the best factor */
2234+
if (size <= (PY_SSIZE_T_MAX - size / 2)) {
2235+
size += size / 2;
2236+
}
2237+
#else
2238+
/* On Linux, overallocate by 25% is the best factor */
2239+
if (size <= (PY_SSIZE_T_MAX - size / 4)) {
2240+
size += size / 4;
2241+
}
2242+
#endif
2243+
}
2244+
2245+
if (writer->obj != NULL) {
2246+
if (_PyBytes_Resize(&writer->obj, size)) {
2247+
return -1;
2248+
}
2249+
assert(writer->obj != NULL);
2250+
}
2251+
else {
2252+
writer->obj = PyBytes_FromStringAndSize(NULL, size);
2253+
if (writer->obj == NULL) {
2254+
return -1;
2255+
}
2256+
assert((size_t)size > sizeof(writer->small_buffer));
2257+
memcpy(PyBytes_AS_STRING(writer->obj),
2258+
writer->small_buffer,
2259+
sizeof(writer->small_buffer));
2260+
}
2261+
return 0;
2262+
}
2263+
2264+
static inline void*
2265+
PyBytesWriter_GetData(PyBytesWriter *writer)
2266+
{
2267+
if (writer->obj == NULL) {
2268+
return writer->small_buffer;
2269+
}
2270+
else {
2271+
return PyBytes_AS_STRING(writer->obj);
2272+
}
2273+
}
2274+
2275+
static inline Py_ssize_t
2276+
PyBytesWriter_GetSize(PyBytesWriter *writer)
2277+
{
2278+
return writer->size;
2279+
}
2280+
2281+
static inline void
2282+
PyBytesWriter_Discard(PyBytesWriter *writer)
2283+
{
2284+
if (writer == NULL) {
2285+
return;
2286+
}
2287+
2288+
Py_XDECREF(writer->obj);
2289+
PyMem_Free(writer);
2290+
}
2291+
2292+
static inline PyBytesWriter*
2293+
PyBytesWriter_Create(Py_ssize_t size)
2294+
{
2295+
if (size < 0) {
2296+
PyErr_SetString(PyExc_ValueError, "size must be >= 0");
2297+
return NULL;
2298+
}
2299+
2300+
PyBytesWriter *writer = (PyBytesWriter*)PyMem_Malloc(sizeof(PyBytesWriter));
2301+
if (writer == NULL) {
2302+
PyErr_NoMemory();
2303+
return NULL;
2304+
}
2305+
2306+
writer->obj = NULL;
2307+
writer->size = 0;
2308+
2309+
if (size >= 1) {
2310+
if (_PyBytesWriter_Resize_impl(writer, size, 0) < 0) {
2311+
PyBytesWriter_Discard(writer);
2312+
return NULL;
2313+
}
2314+
writer->size = size;
2315+
}
2316+
return writer;
2317+
}
2318+
2319+
static inline PyObject*
2320+
PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
2321+
{
2322+
PyObject *result;
2323+
if (size == 0) {
2324+
result = PyBytes_FromStringAndSize("", 0);
2325+
}
2326+
else if (writer->obj != NULL) {
2327+
if (size != PyBytes_GET_SIZE(writer->obj)) {
2328+
if (_PyBytes_Resize(&writer->obj, size)) {
2329+
goto error;
2330+
}
2331+
}
2332+
result = writer->obj;
2333+
writer->obj = NULL;
2334+
}
2335+
else {
2336+
result = PyBytes_FromStringAndSize(writer->small_buffer, size);
2337+
}
2338+
PyBytesWriter_Discard(writer);
2339+
return result;
2340+
2341+
error:
2342+
PyBytesWriter_Discard(writer);
2343+
return NULL;
2344+
}
2345+
2346+
static inline PyObject*
2347+
PyBytesWriter_Finish(PyBytesWriter *writer)
2348+
{
2349+
return PyBytesWriter_FinishWithSize(writer, writer->size);
2350+
}
2351+
2352+
static inline PyObject*
2353+
PyBytesWriter_FinishWithPointer(PyBytesWriter *writer, void *buf)
2354+
{
2355+
Py_ssize_t size = (char*)buf - (char*)PyBytesWriter_GetData(writer);
2356+
if (size < 0 || size > _PyBytesWriter_GetAllocated(writer)) {
2357+
PyBytesWriter_Discard(writer);
2358+
PyErr_SetString(PyExc_ValueError, "invalid end pointer");
2359+
return NULL;
2360+
}
2361+
2362+
return PyBytesWriter_FinishWithSize(writer, size);
2363+
}
2364+
2365+
static inline int
2366+
PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t size)
2367+
{
2368+
if (size < 0) {
2369+
PyErr_SetString(PyExc_ValueError, "size must be >= 0");
2370+
return -1;
2371+
}
2372+
if (_PyBytesWriter_Resize_impl(writer, size, 1) < 0) {
2373+
return -1;
2374+
}
2375+
writer->size = size;
2376+
return 0;
2377+
}
2378+
2379+
static inline int
2380+
PyBytesWriter_Grow(PyBytesWriter *writer, Py_ssize_t size)
2381+
{
2382+
if (size < 0 && writer->size + size < 0) {
2383+
PyErr_SetString(PyExc_ValueError, "invalid size");
2384+
return -1;
2385+
}
2386+
if (size > PY_SSIZE_T_MAX - writer->size) {
2387+
PyErr_NoMemory();
2388+
return -1;
2389+
}
2390+
size = writer->size + size;
2391+
2392+
if (_PyBytesWriter_Resize_impl(writer, size, 1) < 0) {
2393+
return -1;
2394+
}
2395+
writer->size = size;
2396+
return 0;
2397+
}
2398+
2399+
static inline void*
2400+
PyBytesWriter_GrowAndUpdatePointer(PyBytesWriter *writer,
2401+
Py_ssize_t size, void *buf)
2402+
{
2403+
Py_ssize_t pos = (char*)buf - (char*)PyBytesWriter_GetData(writer);
2404+
if (PyBytesWriter_Grow(writer, size) < 0) {
2405+
return NULL;
2406+
}
2407+
return (char*)PyBytesWriter_GetData(writer) + pos;
2408+
}
2409+
2410+
static inline int
2411+
PyBytesWriter_WriteBytes(PyBytesWriter *writer,
2412+
const void *bytes, Py_ssize_t size)
2413+
{
2414+
if (size < 0) {
2415+
size_t len = strlen((const char*)bytes);
2416+
if (len > (size_t)PY_SSIZE_T_MAX) {
2417+
PyErr_NoMemory();
2418+
return -1;
2419+
}
2420+
size = (Py_ssize_t)len;
2421+
}
2422+
2423+
Py_ssize_t pos = writer->size;
2424+
if (PyBytesWriter_Grow(writer, size) < 0) {
2425+
return -1;
2426+
}
2427+
char *buf = (char*)PyBytesWriter_GetData(writer);
2428+
memcpy(buf + pos, bytes, (size_t)size);
2429+
return 0;
2430+
}
2431+
2432+
static inline int
2433+
PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
2434+
{
2435+
va_list vargs;
2436+
va_start(vargs, format);
2437+
PyObject *str = PyBytes_FromFormatV(format, vargs);
2438+
va_end(vargs);
2439+
2440+
if (str == NULL) {
2441+
return -1;
2442+
}
2443+
int res = PyBytesWriter_WriteBytes(writer,
2444+
PyBytes_AS_STRING(str),
2445+
PyBytes_GET_SIZE(str));
2446+
Py_DECREF(str);
2447+
return res;
2448+
}
2449+
#endif // PY_VERSION_HEX < 0x030E00A7
2450+
2451+
22022452
#ifdef __cplusplus
22032453
}
22042454
#endif

‎tests/test_pythoncapi_compat_cext.c

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2181,6 +2181,103 @@ test_config(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
21812181
#endif
21822182

21832183

2184+
static int
2185+
test_byteswriter_highlevel(void)
2186+
{
2187+
PyObject *obj;
2188+
PyBytesWriter *writer = PyBytesWriter_Create(0);
2189+
if (writer == NULL) {
2190+
goto error;
2191+
}
2192+
if (PyBytesWriter_WriteBytes(writer, "Hello", -1) < 0) {
2193+
goto error;
2194+
}
2195+
if (PyBytesWriter_Format(writer, " %s!", "World") < 0) {
2196+
goto error;
2197+
}
2198+
2199+
obj = PyBytesWriter_Finish(writer);
2200+
if (obj == NULL) {
2201+
return -1;
2202+
}
2203+
assert(PyBytes_Check(obj));
2204+
assert(strcmp(PyBytes_AS_STRING(obj), "Hello World!") == 0);
2205+
Py_DECREF(obj);
2206+
return 0;
2207+
2208+
error:
2209+
PyBytesWriter_Discard(writer);
2210+
return -1;
2211+
}
2212+
2213+
static int
2214+
test_byteswriter_abc(void)
2215+
{
2216+
PyBytesWriter *writer = PyBytesWriter_Create(3);
2217+
if (writer == NULL) {
2218+
return -1;
2219+
}
2220+
2221+
char *str = (char*)PyBytesWriter_GetData(writer);
2222+
memcpy(str, "abc", 3);
2223+
2224+
PyObject *obj = PyBytesWriter_Finish(writer);
2225+
if (obj == NULL) {
2226+
return -1;
2227+
}
2228+
assert(PyBytes_Check(obj));
2229+
assert(strcmp(PyBytes_AS_STRING(obj), "abc") == 0);
2230+
Py_DECREF(obj);
2231+
return 0;
2232+
}
2233+
2234+
static int
2235+
test_byteswriter_grow(void)
2236+
{
2237+
PyBytesWriter *writer = PyBytesWriter_Create(10);
2238+
if (writer == NULL) {
2239+
return -1;
2240+
}
2241+
2242+
char *buf = (char*)PyBytesWriter_GetData(writer);
2243+
memcpy(buf, "Hello ", strlen("Hello "));
2244+
buf += strlen("Hello ");
2245+
2246+
buf = (char*)PyBytesWriter_GrowAndUpdatePointer(writer, 10, buf);
2247+
if (buf == NULL) {
2248+
PyBytesWriter_Discard(writer);
2249+
return -1;
2250+
}
2251+
2252+
memcpy(buf, "World", strlen("World"));
2253+
buf += strlen("World");
2254+
2255+
PyObject *obj = PyBytesWriter_FinishWithPointer(writer, buf);
2256+
if (obj == NULL) {
2257+
return -1;
2258+
}
2259+
assert(PyBytes_Check(obj));
2260+
assert(strcmp(PyBytes_AS_STRING(obj), "Hello World") == 0);
2261+
Py_DECREF(obj);
2262+
return 0;
2263+
}
2264+
2265+
static PyObject *
2266+
test_byteswriter(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
2267+
{
2268+
if (test_byteswriter_highlevel() < 0) {
2269+
return NULL;
2270+
}
2271+
if (test_byteswriter_abc() < 0) {
2272+
return NULL;
2273+
}
2274+
if (test_byteswriter_grow() < 0) {
2275+
return NULL;
2276+
}
2277+
Py_RETURN_NONE;
2278+
}
2279+
2280+
21842281
static struct PyMethodDef methods[] = {
21852282
{"test_object", test_object, METH_NOARGS, _Py_NULL},
21862283
{"test_py_is", test_py_is, METH_NOARGS, _Py_NULL},
@@ -2232,6 +2329,7 @@ static struct PyMethodDef methods[] = {
22322329
#if 0x03090000 <= PY_VERSION_HEX && !defined(PYPY_VERSION)
22332330
{"test_config", test_config, METH_NOARGS, _Py_NULL},
22342331
#endif
2332+
{"test_byteswriter", test_byteswriter, METH_NOARGS, _Py_NULL},
22352333
{_Py_NULL, _Py_NULL, 0, _Py_NULL}
22362334
};
22372335

0 commit comments

Comments
 (0)