Newer
Older
#define PY_SSIZE_T_CLEAN
#define Py_LIMITED_API 0x030B0000
#include <stdlib.h>
#include <Python.h>
#include "zlib.h"
#define DEF_MEM_LEVEL 8
static void*
PyZlib_Malloc(voidpf ctx, uInt items, uInt size)
{
if (size != 0 && items > (size_t)PY_SSIZE_T_MAX / size)
return NULL;
/* PyMem_Malloc() cannot be used: the GIL is not held when
inflate() and deflate() are called */
/* The builtin zlib module uses PyMem_RawMalloc here, but this was only
added to the stable ABI from Python 3.13, so use plain malloc for now */
return malloc((size_t)items * (size_t)size);
}
static void
PyZlib_Free(voidpf ctx, void *ptr)
{
/* The builtin zlib module uses PyMem_RawFree here, but this was only
added to the stable ABI from Python 3.13, so use plain free for now */
free(ptr);
}
// zlib_error copied from CPython
static void
zlib_error(z_stream zst, int err, const char *msg)
{
const char *zmsg = Z_NULL;
/* In case of a version mismatch, zst.msg won't be initialized.
Check for this case first, before looking at zst.msg. */
if (err == Z_VERSION_ERROR)
zmsg = "library version mismatch";
if (zmsg == Z_NULL)
zmsg = zst.msg;
if (zmsg == Z_NULL) {
switch (err) {
case Z_BUF_ERROR:
zmsg = "incomplete or truncated stream";
break;
case Z_STREAM_ERROR:
zmsg = "inconsistent stream state";
break;
case Z_DATA_ERROR:
zmsg = "invalid input data";
break;
}
}
if (zmsg == Z_NULL)
PyErr_Format(PyExc_RuntimeError, "Error %d %s", err, msg);
else
PyErr_Format(PyExc_RuntimeError, "Error %d %s: %.200s", err, msg, zmsg);
}
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
static PyObject *
compress_into(PyObject *module, PyObject *args) {
PyObject *return_value = NULL;
Py_buffer input, output;
Py_ssize_t bytes_written;
int level, wbits;
z_stream zst;
if (!PyArg_ParseTuple(args, "y*w*ii", &input, &output, &level, &wbits)) {
return NULL;
}
if (output.len <= 0) {
PyErr_SetString(PyExc_ValueError, "Output buffer may not be 0 size");
goto done;
}
zst.opaque = NULL;
zst.zalloc = PyZlib_Malloc;
zst.zfree = PyZlib_Free;
zst.next_in = input.buf;
zst.avail_in = input.len;
zst.next_out = output.buf;
zst.avail_out = output.len;
int err = deflateInit2(&zst, level, Z_DEFLATED, wbits, DEF_MEM_LEVEL,
Z_DEFAULT_STRATEGY);
switch (err) {
case Z_OK:
break;
case Z_MEM_ERROR:
PyErr_SetString(PyExc_MemoryError,
"Out of memory while compressing data");
goto done;
case Z_STREAM_ERROR:
PyErr_SetString(PyExc_ValueError, "Bad compression level");
goto done;
default:
deflateEnd(&zst);
goto done;
}
Py_BEGIN_ALLOW_THREADS
err = deflate(&zst, Z_FINISH);
Py_END_ALLOW_THREADS
switch (err) {
case Z_STREAM_END:
break;
case Z_OK:
case Z_BUF_ERROR:
deflateEnd(&zst);
PyErr_SetString(PyExc_BufferError, "Not enough space in output buffer");
goto done;
default:
deflateEnd(&zst);
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
goto done;
}
bytes_written = output.len - zst.avail_out;
deflateEnd(&zst);
return_value = PyLong_FromSsize_t(bytes_written);
done:
PyBuffer_Release(&input);
PyBuffer_Release(&output);
return return_value;
}
static PyMethodDef ZlibIntoMethods[] = {
{"compress_into", compress_into, METH_VARARGS, "zlib compress data into a buffer"},
{NULL, NULL, 0, NULL}
};
static struct PyModuleDef zlibintomodule = {
PyModuleDef_HEAD_INIT,
"zlib_into",
NULL, // docstring
-1,
ZlibIntoMethods
};
PyMODINIT_FUNC
PyInit_zlib_into(void) {
return PyModule_Create(&zlibintomodule);
}