mirrors/cpython
1
0
Fork 0
mirror of https://github.com/python/cpython.git synced 2025-09-26 18:29:57 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 10

bpo-41486: Faster bz2/lzma/zlib via new output buffering (GH-21740)

Browse source 
Faster bz2/lzma/zlib via new output buffering.
Also adds .readall() function to _compression.DecompressReader class
to take best advantage of this in the consume-all-output at once scenario.

Often a 5-20% speedup in common scenarios due to less data copying.

Contributed by Ma Lin.
This commit is contained in:
Ma Lin 2021-04-28 14:58:54 +08:00 • committed by GitHub
parent a5e64444e6
commit f9bedb630e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 670 additions and 254 deletions
Download patch file Download diff file Expand all files Collapse all files

6
Doc/whatsnew/3.10.rst
View file

@ -1243,6 +1243,12 @@ Optimizations
for more details. (Contributed by Victor Stinner and Pablo Galindo in for more details. (Contributed by Victor Stinner and Pablo Galindo in
:issue:`38980`.) :issue:`38980`.)
* Use a new output buffer management code for :mod:`bz2` / :mod:`lzma` /
:mod:`zlib` modules, and add ``.readall()`` function to
``_compression.DecompressReader`` class. bz2 decompression 1.09x ~ 1.17x
faster, lzma decompression 1.20x ~ 1.32x faster, ``GzipFile.read(-1)`` 1.11x
~ 1.18x faster. (Contributed by Ma Lin, reviewed by Gregory P. Smith, in :issue:`41486`)
* Function parameters and their annotations are no longer computed at runtime, * Function parameters and their annotations are no longer computed at runtime,
but rather at compilation time. They are stored as a tuple of strings at the but rather at compilation time. They are stored as a tuple of strings at the
bytecode level. It is now around 2 times faster to create a function with bytecode level. It is now around 2 times faster to create a function with

317
Include/internal/pycore_blocks_output_buffer.h Normal file
View file

@ -0,0 +1,317 @@
/*
_BlocksOutputBuffer is used to maintain an output buffer
that has unpredictable size. Suitable for compression/decompression
API (bz2/lzma/zlib) that has stream->next_out and stream->avail_out:
stream->next_out: point to the next output position.
stream->avail_out: the number of available bytes left in the buffer.
It maintains a list of bytes object, so there is no overhead of resizing
the buffer.
Usage:
1, Initialize the struct instance like this:
_BlocksOutputBuffer buffer = {.list = NULL};
Set .list to NULL for _BlocksOutputBuffer_OnError()
2, Initialize the buffer use one of these functions:
_BlocksOutputBuffer_InitAndGrow()
_BlocksOutputBuffer_InitWithSize()
3, If (avail_out == 0), grow the buffer:
_BlocksOutputBuffer_Grow()
4, Get the current outputted data size:
_BlocksOutputBuffer_GetDataSize()
5, Finish the buffer, and return a bytes object:
_BlocksOutputBuffer_Finish()
6, Clean up the buffer when an error occurred:
_BlocksOutputBuffer_OnError()
*/
#ifndef Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H
#define Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H
#ifdef __cplusplus
extern "C" {
#endif
#include "Python.h"
typedef struct {
// List of bytes objects
PyObject *list;
// Number of whole allocated size
Py_ssize_t allocated;
// Max length of the buffer, negative number means unlimited length.
Py_ssize_t max_length;
} _BlocksOutputBuffer;
static const char unable_allocate_msg[] = "Unable to allocate output buffer.";
/* In 32-bit build, the max block size should <= INT32_MAX. */
#define OUTPUT_BUFFER_MAX_BLOCK_SIZE (256*1024*1024)
/* Block size sequence */
#define KB (1024)
#define MB (1024*1024)
const Py_ssize_t BUFFER_BLOCK_SIZE[] =
{ 32*KB, 64*KB, 256*KB, 1*MB, 4*MB, 8*MB, 16*MB, 16*MB,
32*MB, 32*MB, 32*MB, 32*MB, 64*MB, 64*MB, 128*MB, 128*MB,
OUTPUT_BUFFER_MAX_BLOCK_SIZE };
#undef KB
#undef MB
/* According to the block sizes defined by BUFFER_BLOCK_SIZE, the whole
allocated size growth step is:
1 32 KB +32 KB
2 96 KB +64 KB
3 352 KB +256 KB
4 1.34 MB +1 MB
5 5.34 MB +4 MB
6 13.34 MB +8 MB
7 29.34 MB +16 MB
8 45.34 MB +16 MB
9 77.34 MB +32 MB
10 109.34 MB +32 MB
11 141.34 MB +32 MB
12 173.34 MB +32 MB
13 237.34 MB +64 MB
14 301.34 MB +64 MB
15 429.34 MB +128 MB
16 557.34 MB +128 MB
17 813.34 MB +256 MB
18 1069.34 MB +256 MB
19 1325.34 MB +256 MB
20 1581.34 MB +256 MB
21 1837.34 MB +256 MB
22 2093.34 MB +256 MB
...
*/
/* Initialize the buffer, and grow the buffer.
max_length: Max length of the buffer, -1 for unlimited length.
On success, return allocated size (>=0)
On failure, return -1
*/
static inline Py_ssize_t
_BlocksOutputBuffer_InitAndGrow(_BlocksOutputBuffer *buffer,
const Py_ssize_t max_length,
void **next_out)
{
PyObject *b;
Py_ssize_t block_size;
// ensure .list was set to NULL
assert(buffer->list == NULL);
// get block size
if (0 <= max_length && max_length < BUFFER_BLOCK_SIZE[0]) {
block_size = max_length;
} else {
block_size = BUFFER_BLOCK_SIZE[0];
}
// the first block
b = PyBytes_FromStringAndSize(NULL, block_size);
if (b == NULL) {
return -1;
}
// create the list
buffer->list = PyList_New(1);
if (buffer->list == NULL) {
Py_DECREF(b);
return -1;
}
PyList_SET_ITEM(buffer->list, 0, b);
// set variables
buffer->allocated = block_size;
buffer->max_length = max_length;
*next_out = PyBytes_AS_STRING(b);
return block_size;
}
/* Initialize the buffer, with an initial size.
Check block size limit in the outer wrapper function. For example, some libs
accept UINT32_MAX as the maximum block size, then init_size should <= it.
On success, return allocated size (>=0)
On failure, return -1
*/
static inline Py_ssize_t
_BlocksOutputBuffer_InitWithSize(_BlocksOutputBuffer *buffer,
const Py_ssize_t init_size,
void **next_out)
{
PyObject *b;
// ensure .list was set to NULL
assert(buffer->list == NULL);
// the first block
b = PyBytes_FromStringAndSize(NULL, init_size);
if (b == NULL) {
PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
return -1;
}
// create the list
buffer->list = PyList_New(1);
if (buffer->list == NULL) {
Py_DECREF(b);
return -1;
}
PyList_SET_ITEM(buffer->list, 0, b);
// set variables
buffer->allocated = init_size;
buffer->max_length = -1;
*next_out = PyBytes_AS_STRING(b);
return init_size;
}
/* Grow the buffer. The avail_out must be 0, please check it before calling.
On success, return allocated size (>=0)
On failure, return -1
*/
static inline Py_ssize_t
_BlocksOutputBuffer_Grow(_BlocksOutputBuffer *buffer,
void **next_out,
const Py_ssize_t avail_out)
{
PyObject *b;
const Py_ssize_t list_len = Py_SIZE(buffer->list);
Py_ssize_t block_size;
// ensure no gaps in the data
if (avail_out != 0) {
PyErr_SetString(PyExc_SystemError,
"avail_out is non-zero in _BlocksOutputBuffer_Grow().");
return -1;
}
// get block size
if (list_len < (Py_ssize_t) Py_ARRAY_LENGTH(BUFFER_BLOCK_SIZE)) {
block_size = BUFFER_BLOCK_SIZE[list_len];
} else {
block_size = BUFFER_BLOCK_SIZE[Py_ARRAY_LENGTH(BUFFER_BLOCK_SIZE) - 1];
}
// check max_length
if (buffer->max_length >= 0) {
// if (rest == 0), should not grow the buffer.
Py_ssize_t rest = buffer->max_length - buffer->allocated;
assert(rest > 0);
// block_size of the last block
if (block_size > rest) {
block_size = rest;
}
}
// check buffer->allocated overflow
if (block_size > PY_SSIZE_T_MAX - buffer->allocated) {
PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
return -1;
}
// create the block
b = PyBytes_FromStringAndSize(NULL, block_size);
if (b == NULL) {
PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
return -1;
}
if (PyList_Append(buffer->list, b) < 0) {
Py_DECREF(b);
return -1;
}
Py_DECREF(b);
// set variables
buffer->allocated += block_size;
*next_out = PyBytes_AS_STRING(b);
return block_size;
}
/* Return the current outputted data size. */
static inline Py_ssize_t
_BlocksOutputBuffer_GetDataSize(_BlocksOutputBuffer *buffer,
const Py_ssize_t avail_out)
{
return buffer->allocated - avail_out;
}
/* Finish the buffer.
Return a bytes object on success
Return NULL on failure
*/
static inline PyObject *
_BlocksOutputBuffer_Finish(_BlocksOutputBuffer *buffer,
const Py_ssize_t avail_out)
{
PyObject *result, *block;
const Py_ssize_t list_len = Py_SIZE(buffer->list);
// fast path for single block
if ((list_len == 1 && avail_out == 0) ||
(list_len == 2 && Py_SIZE(PyList_GET_ITEM(buffer->list, 1)) == avail_out))
{
block = PyList_GET_ITEM(buffer->list, 0);
Py_INCREF(block);
Py_CLEAR(buffer->list);
return block;
}
// final bytes object
result = PyBytes_FromStringAndSize(NULL, buffer->allocated - avail_out);
if (result == NULL) {
PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
return NULL;
}
// memory copy
if (list_len > 0) {
char *posi = PyBytes_AS_STRING(result);
// blocks except the last one
Py_ssize_t i = 0;
for (; i < list_len-1; i++) {
block = PyList_GET_ITEM(buffer->list, i);
memcpy(posi, PyBytes_AS_STRING(block), Py_SIZE(block));
posi += Py_SIZE(block);
}
// the last block
block = PyList_GET_ITEM(buffer->list, i);
memcpy(posi, PyBytes_AS_STRING(block), Py_SIZE(block) - avail_out);
} else {
assert(Py_SIZE(result) == 0);
}
Py_CLEAR(buffer->list);
return result;
}
/* Clean up the buffer when an error occurred. */
static inline void
_BlocksOutputBuffer_OnError(_BlocksOutputBuffer *buffer)
{
Py_CLEAR(buffer->list);
}
#ifdef __cplusplus
}
#endif
#endif /* Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H */

12
Lib/_compression.py
View file

@ -1,7 +1,7 @@
"""Internal classes used by the gzip, lzma and bz2 modules""" """Internal classes used by the gzip, lzma and bz2 modules"""
import io import io
import sys
BUFFER_SIZE = io.DEFAULT_BUFFER_SIZE # Compressed data read chunk size BUFFER_SIZE = io.DEFAULT_BUFFER_SIZE # Compressed data read chunk size
@ -110,6 +110,16 @@ class DecompressReader(io.RawIOBase):
self._pos += len(data) self._pos += len(data)
return data return data
def readall(self):
chunks = []
# sys.maxsize means the max length of output buffer is unlimited,
# so that the whole input buffer can be decompressed within one
# .decompress() call.
while data := self.read(sys.maxsize):
chunks.append(data)
return b"".join(chunks)
# Rewind the file to the beginning of the data stream. # Rewind the file to the beginning of the data stream.
def _rewind(self): def _rewind(self):
self._fp.seek(0) self._fp.seek(0)

4
Misc/NEWS.d/next/Library/2020-10-16-15-34-30.bpo-41486.Mu9Iit.rst Normal file
View file

@ -0,0 +1,4 @@
Use a new output buffer management code for :mod:`bz2` / :mod:`lzma` /
:mod:`zlib` modules, and add ``.readall()`` function to
``_compression.DecompressReader`` class. These bring some performance
improvements. Patch by Ma Lin.

148
Modules/_bz2module.c
View file

@ -8,6 +8,59 @@
#include <bzlib.h> #include <bzlib.h>
#include <stdio.h> #include <stdio.h>
// Blocks output buffer wrappers
#include "pycore_blocks_output_buffer.h"
#if OUTPUT_BUFFER_MAX_BLOCK_SIZE > UINT32_MAX
#error "The maximum block size accepted by libbzip2 is UINT32_MAX."
#endif
/* On success, return value >= 0
On failure, return -1 */
static inline Py_ssize_t
Buffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
char **next_out, uint32_t *avail_out)
{
Py_ssize_t allocated;
allocated = _BlocksOutputBuffer_InitAndGrow(
buffer, max_length, (void**) next_out);
*avail_out = (uint32_t) allocated;
return allocated;
}
/* On success, return value >= 0
On failure, return -1 */
static inline Py_ssize_t
Buffer_Grow(_BlocksOutputBuffer *buffer,
char **next_out, uint32_t *avail_out)
{
Py_ssize_t allocated;
allocated = _BlocksOutputBuffer_Grow(
buffer, (void**) next_out, (Py_ssize_t) *avail_out);
*avail_out = (uint32_t) allocated;
return allocated;
}
static inline Py_ssize_t
Buffer_GetDataSize(_BlocksOutputBuffer *buffer, uint32_t avail_out)
{
return _BlocksOutputBuffer_GetDataSize(buffer, (Py_ssize_t) avail_out);
}
static inline PyObject *
Buffer_Finish(_BlocksOutputBuffer *buffer, uint32_t avail_out)
{
return _BlocksOutputBuffer_Finish(buffer, (Py_ssize_t) avail_out);
}
static inline void
Buffer_OnError(_BlocksOutputBuffer *buffer)
{
_BlocksOutputBuffer_OnError(buffer);
}
#ifndef BZ_CONFIG_ERROR #ifndef BZ_CONFIG_ERROR
#define BZ2_bzCompress bzCompress #define BZ2_bzCompress bzCompress
@ -115,52 +168,22 @@ catch_bz2_error(int bzerror)
} }
} }
#if BUFSIZ < 8192
#define INITIAL_BUFFER_SIZE 8192
#else
#define INITIAL_BUFFER_SIZE BUFSIZ
#endif
static int
grow_buffer(PyObject **buf, Py_ssize_t max_length)
{
/* Expand the buffer by an amount proportional to the current size,
giving us amortized linear-time behavior. Use a less-than-double
growth factor to avoid excessive allocation. */
size_t size = PyBytes_GET_SIZE(*buf);
size_t new_size = size + (size >> 3) + 6;
if (max_length > 0 && new_size > (size_t) max_length)
new_size = (size_t) max_length;
if (new_size > size) {
return _PyBytes_Resize(buf, new_size);
} else { /* overflow */
PyErr_SetString(PyExc_OverflowError,
"Unable to allocate buffer - output too large");
return -1;
}
}
/* BZ2Compressor class. */ /* BZ2Compressor class. */
static PyObject * static PyObject *
compress(BZ2Compressor *c, char *data, size_t len, int action) compress(BZ2Compressor *c, char *data, size_t len, int action)
{ {
size_t data_size = 0;
PyObject *result; PyObject *result;
_BlocksOutputBuffer buffer = {.list = NULL};
result = PyBytes_FromStringAndSize(NULL, INITIAL_BUFFER_SIZE); if (Buffer_InitAndGrow(&buffer, -1, &c->bzs.next_out, &c->bzs.avail_out) < 0) {
if (result == NULL) goto error;
return NULL; }
c->bzs.next_in = data; c->bzs.next_in = data;
c->bzs.avail_in = 0; c->bzs.avail_in = 0;
c->bzs.next_out = PyBytes_AS_STRING(result);
c->bzs.avail_out = INITIAL_BUFFER_SIZE;
for (;;) { for (;;) {
char *this_out;
int bzerror; int bzerror;
/* On a 64-bit system, len might not fit in avail_in (an unsigned int). /* On a 64-bit system, len might not fit in avail_in (an unsigned int).
@ -175,21 +198,15 @@ compress(BZ2Compressor *c, char *data, size_t len, int action)
break; break;
if (c->bzs.avail_out == 0) { if (c->bzs.avail_out == 0) {
size_t buffer_left = PyBytes_GET_SIZE(result) - data_size; if (Buffer_Grow(&buffer, &c->bzs.next_out, &c->bzs.avail_out) < 0) {
if (buffer_left == 0) {
if (grow_buffer(&result, -1) < 0)
goto error; goto error;
c->bzs.next_out = PyBytes_AS_STRING(result) + data_size;
buffer_left = PyBytes_GET_SIZE(result) - data_size;
} }
c->bzs.avail_out = (unsigned int)Py_MIN(buffer_left, UINT_MAX);
} }
Py_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS
this_out = c->bzs.next_out;
bzerror = BZ2_bzCompress(&c->bzs, action); bzerror = BZ2_bzCompress(&c->bzs, action);
data_size += c->bzs.next_out - this_out;
Py_END_ALLOW_THREADS Py_END_ALLOW_THREADS
if (catch_bz2_error(bzerror)) if (catch_bz2_error(bzerror))
goto error; goto error;
@ -197,13 +214,14 @@ compress(BZ2Compressor *c, char *data, size_t len, int action)
if (action == BZ_FINISH && bzerror == BZ_STREAM_END) if (action == BZ_FINISH && bzerror == BZ_STREAM_END)
break; break;
} }
if (data_size != (size_t)PyBytes_GET_SIZE(result))
if (_PyBytes_Resize(&result, data_size) < 0) result = Buffer_Finish(&buffer, c->bzs.avail_out);
goto error; if (result != NULL) {
return result; return result;
}
error: error:
Py_XDECREF(result); Buffer_OnError(&buffer);
return NULL; return NULL;
} }
@ -420,36 +438,29 @@ decompress_buf(BZ2Decompressor *d, Py_ssize_t max_length)
/* data_size is strictly positive, but because we repeatedly have to /* data_size is strictly positive, but because we repeatedly have to
compare against max_length and PyBytes_GET_SIZE we declare it as compare against max_length and PyBytes_GET_SIZE we declare it as
signed */ signed */
Py_ssize_t data_size = 0;
PyObject *result; PyObject *result;
_BlocksOutputBuffer buffer = {.list = NULL};
bz_stream *bzs = &d->bzs; bz_stream *bzs = &d->bzs;
if (max_length < 0 || max_length >= INITIAL_BUFFER_SIZE) if (Buffer_InitAndGrow(&buffer, max_length, &bzs->next_out, &bzs->avail_out) < 0) {
result = PyBytes_FromStringAndSize(NULL, INITIAL_BUFFER_SIZE); goto error;
else }
result = PyBytes_FromStringAndSize(NULL, max_length);
if (result == NULL)
return NULL;
bzs->next_out = PyBytes_AS_STRING(result);
for (;;) { for (;;) {
int bzret; int bzret;
size_t avail;
/* On a 64-bit system, buffer length might not fit in avail_out, so we /* On a 64-bit system, buffer length might not fit in avail_out, so we
do decompression in chunks of no more than UINT_MAX bytes do decompression in chunks of no more than UINT_MAX bytes
each. Note that the expression for `avail` is guaranteed to be each. Note that the expression for `avail` is guaranteed to be
positive, so the cast is safe. */ positive, so the cast is safe. */
avail = (size_t) (PyBytes_GET_SIZE(result) - data_size);
bzs->avail_out = (unsigned int)Py_MIN(avail, UINT_MAX);
bzs->avail_in = (unsigned int)Py_MIN(d->bzs_avail_in_real, UINT_MAX); bzs->avail_in = (unsigned int)Py_MIN(d->bzs_avail_in_real, UINT_MAX);
d->bzs_avail_in_real -= bzs->avail_in; d->bzs_avail_in_real -= bzs->avail_in;
Py_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS
bzret = BZ2_bzDecompress(bzs); bzret = BZ2_bzDecompress(bzs);
data_size = bzs->next_out - PyBytes_AS_STRING(result);
d->bzs_avail_in_real += bzs->avail_in;
Py_END_ALLOW_THREADS Py_END_ALLOW_THREADS
d->bzs_avail_in_real += bzs->avail_in;
if (catch_bz2_error(bzret)) if (catch_bz2_error(bzret))
goto error; goto error;
if (bzret == BZ_STREAM_END) { if (bzret == BZ_STREAM_END) {
@ -458,22 +469,21 @@ decompress_buf(BZ2Decompressor *d, Py_ssize_t max_length)
} else if (d->bzs_avail_in_real == 0) { } else if (d->bzs_avail_in_real == 0) {
break; break;
} else if (bzs->avail_out == 0) { } else if (bzs->avail_out == 0) {
if (data_size == max_length) if (Buffer_GetDataSize(&buffer, bzs->avail_out) == max_length)
break; break;
if (data_size == PyBytes_GET_SIZE(result) && if (Buffer_Grow(&buffer, &bzs->next_out, &bzs->avail_out) < 0) {
grow_buffer(&result, max_length) == -1)
goto error; goto error;
bzs->next_out = PyBytes_AS_STRING(result) + data_size;
} }
} }
if (data_size != PyBytes_GET_SIZE(result)) }
if (_PyBytes_Resize(&result, data_size) == -1)
goto error;
result = Buffer_Finish(&buffer, bzs->avail_out);
if (result != NULL) {
return result; return result;
}
error: error:
Py_XDECREF(result); Buffer_OnError(&buffer);
return NULL; return NULL;
} }

142
Modules/_lzmamodule.c
View file

@ -15,6 +15,60 @@
#include <lzma.h> #include <lzma.h>
// Blocks output buffer wrappers
#include "pycore_blocks_output_buffer.h"
#if OUTPUT_BUFFER_MAX_BLOCK_SIZE > SIZE_MAX
#error "The maximum block size accepted by liblzma is SIZE_MAX."
#endif
/* On success, return value >= 0
On failure, return -1 */
static inline Py_ssize_t
Buffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
uint8_t **next_out, size_t *avail_out)
{
Py_ssize_t allocated;
allocated = _BlocksOutputBuffer_InitAndGrow(
buffer, max_length, (void**) next_out);
*avail_out = (size_t) allocated;
return allocated;
}
/* On success, return value >= 0
On failure, return -1 */
static inline Py_ssize_t
Buffer_Grow(_BlocksOutputBuffer *buffer,
uint8_t **next_out, size_t *avail_out)
{
Py_ssize_t allocated;
allocated = _BlocksOutputBuffer_Grow(
buffer, (void**) next_out, (Py_ssize_t) *avail_out);
*avail_out = (size_t) allocated;
return allocated;
}
static inline Py_ssize_t
Buffer_GetDataSize(_BlocksOutputBuffer *buffer, size_t avail_out)
{
return _BlocksOutputBuffer_GetDataSize(buffer, (Py_ssize_t) avail_out);
}
static inline PyObject *
Buffer_Finish(_BlocksOutputBuffer *buffer, size_t avail_out)
{
return _BlocksOutputBuffer_Finish(buffer, (Py_ssize_t) avail_out);
}
static inline void
Buffer_OnError(_BlocksOutputBuffer *buffer)
{
_BlocksOutputBuffer_OnError(buffer);
}
#define ACQUIRE_LOCK(obj) do { \ #define ACQUIRE_LOCK(obj) do { \
if (!PyThread_acquire_lock((obj)->lock, 0)) { \ if (!PyThread_acquire_lock((obj)->lock, 0)) { \
Py_BEGIN_ALLOW_THREADS \ Py_BEGIN_ALLOW_THREADS \
@ -128,25 +182,6 @@ PyLzma_Free(void *opaque, void *ptr)
PyMem_RawFree(ptr); PyMem_RawFree(ptr);
} }
#if BUFSIZ < 8192
#define INITIAL_BUFFER_SIZE 8192
#else
#define INITIAL_BUFFER_SIZE BUFSIZ
#endif
static int
grow_buffer(PyObject **buf, Py_ssize_t max_length)
{
Py_ssize_t size = PyBytes_GET_SIZE(*buf);
Py_ssize_t newsize = size + (size >> 3) + 6;
if (max_length > 0 && newsize > max_length) {
newsize = max_length;
}
return _PyBytes_Resize(buf, newsize);
}
/* Some custom type conversions for PyArg_ParseTupleAndKeywords(), /* Some custom type conversions for PyArg_ParseTupleAndKeywords(),
since the predefined conversion specifiers do not suit our needs: since the predefined conversion specifiers do not suit our needs:
@ -510,29 +545,27 @@ class lzma_filter_converter(CConverter):
static PyObject * static PyObject *
compress(Compressor *c, uint8_t *data, size_t len, lzma_action action) compress(Compressor *c, uint8_t *data, size_t len, lzma_action action)
{ {
Py_ssize_t data_size = 0;
PyObject *result; PyObject *result;
_BlocksOutputBuffer buffer = {.list = NULL};
_lzma_state *state = PyType_GetModuleState(Py_TYPE(c)); _lzma_state *state = PyType_GetModuleState(Py_TYPE(c));
assert(state != NULL); assert(state != NULL);
result = PyBytes_FromStringAndSize(NULL, INITIAL_BUFFER_SIZE); if (Buffer_InitAndGrow(&buffer, -1, &c->lzs.next_out, &c->lzs.avail_out) < 0) {
if (result == NULL) { goto error;
return NULL;
} }
c->lzs.next_in = data; c->lzs.next_in = data;
c->lzs.avail_in = len; c->lzs.avail_in = len;
c->lzs.next_out = (uint8_t *)PyBytes_AS_STRING(result);
c->lzs.avail_out = PyBytes_GET_SIZE(result);
for (;;) { for (;;) {
lzma_ret lzret; lzma_ret lzret;
Py_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS
lzret = lzma_code(&c->lzs, action); lzret = lzma_code(&c->lzs, action);
data_size = (char *)c->lzs.next_out - PyBytes_AS_STRING(result); Py_END_ALLOW_THREADS
if (lzret == LZMA_BUF_ERROR && len == 0 && c->lzs.avail_out > 0) { if (lzret == LZMA_BUF_ERROR && len == 0 && c->lzs.avail_out > 0) {
lzret = LZMA_OK; /* That wasn't a real error */ lzret = LZMA_OK; /* That wasn't a real error */
} }
Py_END_ALLOW_THREADS
if (catch_lzma_error(state, lzret)) { if (catch_lzma_error(state, lzret)) {
goto error; goto error;
} }
@ -540,20 +573,19 @@ compress(Compressor *c, uint8_t *data, size_t len, lzma_action action)
(action == LZMA_FINISH && lzret == LZMA_STREAM_END)) { (action == LZMA_FINISH && lzret == LZMA_STREAM_END)) {
break; break;
} else if (c->lzs.avail_out == 0) { } else if (c->lzs.avail_out == 0) {
if (grow_buffer(&result, -1) == -1) if (Buffer_Grow(&buffer, &c->lzs.next_out, &c->lzs.avail_out) < 0) {
goto error;
c->lzs.next_out = (uint8_t *)PyBytes_AS_STRING(result) + data_size;
c->lzs.avail_out = PyBytes_GET_SIZE(result) - data_size;
}
}
if (data_size != PyBytes_GET_SIZE(result))
if (_PyBytes_Resize(&result, data_size) == -1) {
goto error; goto error;
} }
}
}
result = Buffer_Finish(&buffer, c->lzs.avail_out);
if (result != NULL) {
return result; return result;
}
error: error:
Py_XDECREF(result); Buffer_OnError(&buffer);
return NULL; return NULL;
} }
@ -896,36 +928,26 @@ static PyType_Spec lzma_compressor_type_spec = {
static PyObject* static PyObject*
decompress_buf(Decompressor *d, Py_ssize_t max_length) decompress_buf(Decompressor *d, Py_ssize_t max_length)
{ {
Py_ssize_t data_size = 0;
PyObject *result; PyObject *result;
lzma_stream *lzs = &d->lzs; lzma_stream *lzs = &d->lzs;
_BlocksOutputBuffer buffer = {.list = NULL};
_lzma_state *state = PyType_GetModuleState(Py_TYPE(d)); _lzma_state *state = PyType_GetModuleState(Py_TYPE(d));
assert(state != NULL); assert(state != NULL);
if (max_length < 0 || max_length >= INITIAL_BUFFER_SIZE) { if (Buffer_InitAndGrow(&buffer, max_length, &lzs->next_out, &lzs->avail_out) < 0) {
result = PyBytes_FromStringAndSize(NULL, INITIAL_BUFFER_SIZE); goto error;
} }
else {
result = PyBytes_FromStringAndSize(NULL, max_length);
}
if (result == NULL) {
return NULL;
}
lzs->next_out = (uint8_t *)PyBytes_AS_STRING(result);
lzs->avail_out = PyBytes_GET_SIZE(result);
for (;;) { for (;;) {
lzma_ret lzret; lzma_ret lzret;
Py_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS
lzret = lzma_code(lzs, LZMA_RUN); lzret = lzma_code(lzs, LZMA_RUN);
data_size = (char *)lzs->next_out - PyBytes_AS_STRING(result); Py_END_ALLOW_THREADS
if (lzret == LZMA_BUF_ERROR && lzs->avail_in == 0 && lzs->avail_out > 0) { if (lzret == LZMA_BUF_ERROR && lzs->avail_in == 0 && lzs->avail_out > 0) {
lzret = LZMA_OK; /* That wasn't a real error */ lzret = LZMA_OK; /* That wasn't a real error */
} }
Py_END_ALLOW_THREADS
if (catch_lzma_error(state, lzret)) { if (catch_lzma_error(state, lzret)) {
goto error; goto error;
} }
@ -940,28 +962,24 @@ decompress_buf(Decompressor *d, Py_ssize_t max_length)
Maybe lzs's internal state still have a few bytes Maybe lzs's internal state still have a few bytes
can be output, grow the output buffer and continue can be output, grow the output buffer and continue
if max_lengh < 0. */ if max_lengh < 0. */
if (data_size == max_length) { if (Buffer_GetDataSize(&buffer, lzs->avail_out) == max_length) {
break; break;
} }
if (grow_buffer(&result, max_length) == -1) { if (Buffer_Grow(&buffer, &lzs->next_out, &lzs->avail_out) < 0) {
goto error; goto error;
} }
lzs->next_out = (uint8_t *)PyBytes_AS_STRING(result) + data_size;
lzs->avail_out = PyBytes_GET_SIZE(result) - data_size;
} else if (lzs->avail_in == 0) { } else if (lzs->avail_in == 0) {
break; break;
} }
} }
if (data_size != PyBytes_GET_SIZE(result)) {
if (_PyBytes_Resize(&result, data_size) == -1) {
goto error;
}
}
result = Buffer_Finish(&buffer, lzs->avail_out);
if (result != NULL) {
return result; return result;
}
error: error:
Py_XDECREF(result); Buffer_OnError(&buffer);
return NULL; return NULL;
} }
@ -1042,7 +1060,7 @@ decompress(Decompressor *d, uint8_t *data, size_t len, Py_ssize_t max_length)
be output, try to output them next time. */ be output, try to output them next time. */
d->needs_input = 0; d->needs_input = 0;
/* if max_length < 0, lzs->avail_out always > 0 */ /* If max_length < 0, lzs->avail_out always > 0 */
assert(max_length >= 0); assert(max_length >= 0);
} else { } else {
/* Input buffer exhausted, output buffer has space. */ /* Input buffer exhausted, output buffer has space. */

267
Modules/zlibmodule.c
View file

@ -9,6 +9,79 @@
#include "structmember.h" // PyMemberDef #include "structmember.h" // PyMemberDef
#include "zlib.h" #include "zlib.h"
// Blocks output buffer wrappers
#include "pycore_blocks_output_buffer.h"
#if OUTPUT_BUFFER_MAX_BLOCK_SIZE > UINT32_MAX
#error "The maximum block size accepted by zlib is UINT32_MAX."
#endif
/* On success, return value >= 0
On failure, return -1 */
static inline Py_ssize_t
Buffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
Bytef **next_out, uint32_t *avail_out)
{
Py_ssize_t allocated;
allocated = _BlocksOutputBuffer_InitAndGrow(
buffer, max_length, (void**) next_out);
*avail_out = (uint32_t) allocated;
return allocated;
}
/* On success, return value >= 0
On failure, return -1 */
static inline Py_ssize_t
Buffer_InitWithSize(_BlocksOutputBuffer *buffer, Py_ssize_t init_size,
Bytef **next_out, uint32_t *avail_out)
{
Py_ssize_t allocated;
if (init_size < 0 || (size_t)init_size > UINT32_MAX) {
PyErr_SetString(PyExc_ValueError,
"Initial buffer size should (0 <= size <= UINT32_MAX)");
return -1;
}
allocated = _BlocksOutputBuffer_InitWithSize(
buffer, init_size, (void**) next_out);
*avail_out = (uint32_t) allocated;
return allocated;
}
/* On success, return value >= 0
On failure, return -1 */
static inline Py_ssize_t
Buffer_Grow(_BlocksOutputBuffer *buffer,
Bytef **next_out, uint32_t *avail_out)
{
Py_ssize_t allocated;
allocated = _BlocksOutputBuffer_Grow(
buffer, (void**) next_out, (Py_ssize_t) *avail_out);
*avail_out = (uint32_t) allocated;
return allocated;
}
static inline Py_ssize_t
Buffer_GetDataSize(_BlocksOutputBuffer *buffer, uint32_t avail_out)
{
return _BlocksOutputBuffer_GetDataSize(buffer, (Py_ssize_t) avail_out);
}
static inline PyObject *
Buffer_Finish(_BlocksOutputBuffer *buffer, uint32_t avail_out)
{
return _BlocksOutputBuffer_Finish(buffer, (Py_ssize_t) avail_out);
}
static inline void
Buffer_OnError(_BlocksOutputBuffer *buffer)
{
_BlocksOutputBuffer_OnError(buffer);
}
#define ENTER_ZLIB(obj) do { \ #define ENTER_ZLIB(obj) do { \
if (!PyThread_acquire_lock((obj)->lock, 0)) { \ if (!PyThread_acquire_lock((obj)->lock, 0)) { \
@ -149,56 +222,6 @@ arrange_input_buffer(z_stream *zst, Py_ssize_t *remains)
*remains -= zst->avail_in; *remains -= zst->avail_in;
} }
static Py_ssize_t
arrange_output_buffer_with_maximum(z_stream *zst, PyObject **buffer,
Py_ssize_t length,
Py_ssize_t max_length)
{
Py_ssize_t occupied;
if (*buffer == NULL) {
if (!(*buffer = PyBytes_FromStringAndSize(NULL, length)))
return -1;
occupied = 0;
}
else {
occupied = zst->next_out - (Byte *)PyBytes_AS_STRING(*buffer);
if (length == occupied) {
Py_ssize_t new_length;
assert(length <= max_length);
/* can not scale the buffer over max_length */
if (length == max_length)
return -2;
if (length <= (max_length >> 1))
new_length = length << 1;
else
new_length = max_length;
if (_PyBytes_Resize(buffer, new_length) < 0)
return -1;
length = new_length;
}
}
zst->avail_out = (uInt)Py_MIN((size_t)(length - occupied), UINT_MAX);
zst->next_out = (Byte *)PyBytes_AS_STRING(*buffer) + occupied;
return length;
}
static Py_ssize_t
arrange_output_buffer(z_stream *zst, PyObject **buffer, Py_ssize_t length)
{
Py_ssize_t ret;
ret = arrange_output_buffer_with_maximum(zst, buffer, length,
PY_SSIZE_T_MAX);
if (ret == -2)
PyErr_NoMemory();
return ret;
}
/*[clinic input] /*[clinic input]
zlib.compress zlib.compress
@ -215,16 +238,20 @@ static PyObject *
zlib_compress_impl(PyObject *module, Py_buffer *data, int level) zlib_compress_impl(PyObject *module, Py_buffer *data, int level)
/*[clinic end generated code: output=d80906d73f6294c8 input=638d54b6315dbed3]*/ /*[clinic end generated code: output=d80906d73f6294c8 input=638d54b6315dbed3]*/
{ {
PyObject *RetVal = NULL; PyObject *RetVal;
Py_ssize_t obuflen = DEF_BUF_SIZE;
int flush; int flush;
z_stream zst; z_stream zst;
_BlocksOutputBuffer buffer = {.list = NULL};
zlibstate *state = get_zlib_state(module); zlibstate *state = get_zlib_state(module);
Byte *ibuf = data->buf; Byte *ibuf = data->buf;
Py_ssize_t ibuflen = data->len; Py_ssize_t ibuflen = data->len;
if (Buffer_InitAndGrow(&buffer, -1, &zst.next_out, &zst.avail_out) < 0) {
goto error;
}
zst.opaque = NULL; zst.opaque = NULL;
zst.zalloc = PyZlib_Malloc; zst.zalloc = PyZlib_Malloc;
zst.zfree = PyZlib_Free; zst.zfree = PyZlib_Free;
@ -252,11 +279,12 @@ zlib_compress_impl(PyObject *module, Py_buffer *data, int level)
flush = ibuflen == 0 ? Z_FINISH : Z_NO_FLUSH; flush = ibuflen == 0 ? Z_FINISH : Z_NO_FLUSH;
do { do {
obuflen = arrange_output_buffer(&zst, &RetVal, obuflen); if (zst.avail_out == 0) {
if (obuflen < 0) { if (Buffer_Grow(&buffer, &zst.next_out, &zst.avail_out) < 0) {
deflateEnd(&zst); deflateEnd(&zst);
goto error; goto error;
} }
}
Py_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS
err = deflate(&zst, flush); err = deflate(&zst, flush);
@ -276,15 +304,16 @@ zlib_compress_impl(PyObject *module, Py_buffer *data, int level)
err = deflateEnd(&zst); err = deflateEnd(&zst);
if (err == Z_OK) { if (err == Z_OK) {
if (_PyBytes_Resize(&RetVal, zst.next_out - RetVal = Buffer_Finish(&buffer, zst.avail_out);
(Byte *)PyBytes_AS_STRING(RetVal)) < 0) if (RetVal == NULL) {
goto error; goto error;
}
return RetVal; return RetVal;
} }
else else
zlib_error(state, zst, err, "while finishing compression"); zlib_error(state, zst, err, "while finishing compression");
error: error:
Py_XDECREF(RetVal); Buffer_OnError(&buffer);
return NULL; return NULL;
} }
@ -307,11 +336,12 @@ zlib_decompress_impl(PyObject *module, Py_buffer *data, int wbits,
Py_ssize_t bufsize) Py_ssize_t bufsize)
/*[clinic end generated code: output=77c7e35111dc8c42 input=a9ac17beff1f893f]*/ /*[clinic end generated code: output=77c7e35111dc8c42 input=a9ac17beff1f893f]*/
{ {
PyObject *RetVal = NULL; PyObject *RetVal;
Byte *ibuf; Byte *ibuf;
Py_ssize_t ibuflen; Py_ssize_t ibuflen;
int err, flush; int err, flush;
z_stream zst; z_stream zst;
_BlocksOutputBuffer buffer = {.list = NULL};
zlibstate *state = get_zlib_state(module); zlibstate *state = get_zlib_state(module);
@ -322,6 +352,10 @@ zlib_decompress_impl(PyObject *module, Py_buffer *data, int wbits,
bufsize = 1; bufsize = 1;
} }
if (Buffer_InitWithSize(&buffer, bufsize, &zst.next_out, &zst.avail_out) < 0) {
goto error;
}
ibuf = data->buf; ibuf = data->buf;
ibuflen = data->len; ibuflen = data->len;
@ -350,11 +384,12 @@ zlib_decompress_impl(PyObject *module, Py_buffer *data, int wbits,
flush = ibuflen == 0 ? Z_FINISH : Z_NO_FLUSH; flush = ibuflen == 0 ? Z_FINISH : Z_NO_FLUSH;
do { do {
bufsize = arrange_output_buffer(&zst, &RetVal, bufsize); if (zst.avail_out == 0) {
if (bufsize < 0) { if (Buffer_Grow(&buffer, &zst.next_out, &zst.avail_out) < 0) {
inflateEnd(&zst); inflateEnd(&zst);
goto error; goto error;
} }
}
Py_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS
err = inflate(&zst, flush); err = inflate(&zst, flush);
@ -393,14 +428,13 @@ zlib_decompress_impl(PyObject *module, Py_buffer *data, int wbits,
goto error; goto error;
} }
if (_PyBytes_Resize(&RetVal, zst.next_out - RetVal = Buffer_Finish(&buffer, zst.avail_out);
(Byte *)PyBytes_AS_STRING(RetVal)) < 0) if (RetVal != NULL) {
goto error;
return RetVal; return RetVal;
}
error: error:
Py_XDECREF(RetVal); Buffer_OnError(&buffer);
return NULL; return NULL;
} }
@ -633,9 +667,9 @@ zlib_Compress_compress_impl(compobject *self, PyTypeObject *cls,
Py_buffer *data) Py_buffer *data)
/*[clinic end generated code: output=6731b3f0ff357ca6 input=04d00f65ab01d260]*/ /*[clinic end generated code: output=6731b3f0ff357ca6 input=04d00f65ab01d260]*/
{ {
PyObject *RetVal = NULL; PyObject *RetVal;
Py_ssize_t obuflen = DEF_BUF_SIZE;
int err; int err;
_BlocksOutputBuffer buffer = {.list = NULL};
zlibstate *state = PyType_GetModuleState(cls); zlibstate *state = PyType_GetModuleState(cls);
ENTER_ZLIB(self); ENTER_ZLIB(self);
@ -643,13 +677,18 @@ zlib_Compress_compress_impl(compobject *self, PyTypeObject *cls,
self->zst.next_in = data->buf; self->zst.next_in = data->buf;
Py_ssize_t ibuflen = data->len; Py_ssize_t ibuflen = data->len;
if (Buffer_InitAndGrow(&buffer, -1, &self->zst.next_out, &self->zst.avail_out) < 0) {
goto error;
}
do { do {
arrange_input_buffer(&self->zst, &ibuflen); arrange_input_buffer(&self->zst, &ibuflen);
do { do {
obuflen = arrange_output_buffer(&self->zst, &RetVal, obuflen); if (self->zst.avail_out == 0) {
if (obuflen < 0) if (Buffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0)
goto error; goto error;
}
Py_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS
err = deflate(&self->zst, Z_NO_FLUSH); err = deflate(&self->zst, Z_NO_FLUSH);
@ -665,12 +704,14 @@ zlib_Compress_compress_impl(compobject *self, PyTypeObject *cls,
} while (ibuflen != 0); } while (ibuflen != 0);
if (_PyBytes_Resize(&RetVal, self->zst.next_out - RetVal = Buffer_Finish(&buffer, self->zst.avail_out);
(Byte *)PyBytes_AS_STRING(RetVal)) == 0) if (RetVal != NULL) {
goto success; goto success;
}
error: error:
Py_CLEAR(RetVal); Buffer_OnError(&buffer);
RetVal = NULL;
success: success:
LEAVE_ZLIB(self); LEAVE_ZLIB(self);
return RetVal; return RetVal;
@ -746,8 +787,9 @@ zlib_Decompress_decompress_impl(compobject *self, PyTypeObject *cls,
/*[clinic end generated code: output=b024a93c2c922d57 input=bfb37b3864cfb606]*/ /*[clinic end generated code: output=b024a93c2c922d57 input=bfb37b3864cfb606]*/
{ {
int err = Z_OK; int err = Z_OK;
Py_ssize_t ibuflen, obuflen = DEF_BUF_SIZE, hard_limit; Py_ssize_t ibuflen;
PyObject *RetVal = NULL; PyObject *RetVal;
_BlocksOutputBuffer buffer = {.list = NULL};
PyObject *module = PyType_GetModule(cls); PyObject *module = PyType_GetModule(cls);
if (module == NULL) if (module == NULL)
@ -758,34 +800,29 @@ zlib_Decompress_decompress_impl(compobject *self, PyTypeObject *cls,
PyErr_SetString(PyExc_ValueError, "max_length must be non-negative"); PyErr_SetString(PyExc_ValueError, "max_length must be non-negative");
return NULL; return NULL;
} else if (max_length == 0) } else if (max_length == 0)
hard_limit = PY_SSIZE_T_MAX; max_length = -1;
else
hard_limit = max_length;
ENTER_ZLIB(self); ENTER_ZLIB(self);
self->zst.next_in = data->buf; self->zst.next_in = data->buf;
ibuflen = data->len; ibuflen = data->len;
/* limit amount of data allocated to max_length */ if (Buffer_InitAndGrow(&buffer, max_length, &self->zst.next_out, &self->zst.avail_out) < 0) {
if (max_length && obuflen > max_length) goto abort;
obuflen = max_length; }
do { do {
arrange_input_buffer(&self->zst, &ibuflen); arrange_input_buffer(&self->zst, &ibuflen);
do { do {
obuflen = arrange_output_buffer_with_maximum(&self->zst, &RetVal, if (self->zst.avail_out == 0) {
obuflen, hard_limit); if (Buffer_GetDataSize(&buffer, self->zst.avail_out) == max_length) {
if (obuflen == -2) {
if (max_length > 0) {
goto save; goto save;
} }
PyErr_NoMemory(); if (Buffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0) {
}
if (obuflen < 0) {
goto abort; goto abort;
} }
}
Py_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS
err = inflate(&self->zst, Z_SYNC_FLUSH); err = inflate(&self->zst, Z_SYNC_FLUSH);
@ -828,12 +865,14 @@ zlib_Decompress_decompress_impl(compobject *self, PyTypeObject *cls,
goto abort; goto abort;
} }
if (_PyBytes_Resize(&RetVal, self->zst.next_out - RetVal = Buffer_Finish(&buffer, self->zst.avail_out);
(Byte *)PyBytes_AS_STRING(RetVal)) == 0) if (RetVal != NULL) {
goto success; goto success;
}
abort: abort:
Py_CLEAR(RetVal); Buffer_OnError(&buffer);
RetVal = NULL;
success: success:
LEAVE_ZLIB(self); LEAVE_ZLIB(self);
return RetVal; return RetVal;
@ -858,8 +897,8 @@ zlib_Compress_flush_impl(compobject *self, PyTypeObject *cls, int mode)
/*[clinic end generated code: output=c7efd13efd62add2 input=286146e29442eb6c]*/ /*[clinic end generated code: output=c7efd13efd62add2 input=286146e29442eb6c]*/
{ {
int err; int err;
Py_ssize_t length = DEF_BUF_SIZE; PyObject *RetVal;
PyObject *RetVal = NULL; _BlocksOutputBuffer buffer = {.list = NULL};
zlibstate *state = PyType_GetModuleState(cls); zlibstate *state = PyType_GetModuleState(cls);
/* Flushing with Z_NO_FLUSH is a no-op, so there's no point in /* Flushing with Z_NO_FLUSH is a no-op, so there's no point in
@ -872,20 +911,23 @@ zlib_Compress_flush_impl(compobject *self, PyTypeObject *cls, int mode)
self->zst.avail_in = 0; self->zst.avail_in = 0;
do { if (Buffer_InitAndGrow(&buffer, -1, &self->zst.next_out, &self->zst.avail_out) < 0) {
length = arrange_output_buffer(&self->zst, &RetVal, length);
if (length < 0) {
Py_CLEAR(RetVal);
goto error; goto error;
} }
do {
if (self->zst.avail_out == 0) {
if (Buffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0) {
goto error;
}
}
Py_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS
err = deflate(&self->zst, mode); err = deflate(&self->zst, mode);
Py_END_ALLOW_THREADS Py_END_ALLOW_THREADS
if (err == Z_STREAM_ERROR) { if (err == Z_STREAM_ERROR) {
zlib_error(state, self->zst, err, "while flushing"); zlib_error(state, self->zst, err, "while flushing");
Py_CLEAR(RetVal);
goto error; goto error;
} }
} while (self->zst.avail_out == 0); } while (self->zst.avail_out == 0);
@ -898,7 +940,6 @@ zlib_Compress_flush_impl(compobject *self, PyTypeObject *cls, int mode)
err = deflateEnd(&self->zst); err = deflateEnd(&self->zst);
if (err != Z_OK) { if (err != Z_OK) {
zlib_error(state, self->zst, err, "while finishing compression"); zlib_error(state, self->zst, err, "while finishing compression");
Py_CLEAR(RetVal);
goto error; goto error;
} }
else else
@ -910,15 +951,18 @@ zlib_Compress_flush_impl(compobject *self, PyTypeObject *cls, int mode)
*/ */
} else if (err != Z_OK && err != Z_BUF_ERROR) { } else if (err != Z_OK && err != Z_BUF_ERROR) {
zlib_error(state, self->zst, err, "while flushing"); zlib_error(state, self->zst, err, "while flushing");
Py_CLEAR(RetVal);
goto error; goto error;
} }
if (_PyBytes_Resize(&RetVal, self->zst.next_out - RetVal = Buffer_Finish(&buffer, self->zst.avail_out);
(Byte *)PyBytes_AS_STRING(RetVal)) < 0) if (RetVal != NULL) {
Py_CLEAR(RetVal); goto success;
}
error: error:
Buffer_OnError(&buffer);
RetVal = NULL;
success:
LEAVE_ZLIB(self); LEAVE_ZLIB(self);
return RetVal; return RetVal;
} }
@ -1120,8 +1164,9 @@ zlib_Decompress_flush_impl(compobject *self, PyTypeObject *cls,
{ {
int err, flush; int err, flush;
Py_buffer data; Py_buffer data;
PyObject *RetVal = NULL; PyObject *RetVal;
Py_ssize_t ibuflen; Py_ssize_t ibuflen;
_BlocksOutputBuffer buffer = {.list = NULL};
PyObject *module = PyType_GetModule(cls); PyObject *module = PyType_GetModule(cls);
if (module == NULL) { if (module == NULL) {
@ -1144,14 +1189,19 @@ zlib_Decompress_flush_impl(compobject *self, PyTypeObject *cls,
self->zst.next_in = data.buf; self->zst.next_in = data.buf;
ibuflen = data.len; ibuflen = data.len;
if (Buffer_InitWithSize(&buffer, length, &self->zst.next_out, &self->zst.avail_out) < 0) {
goto abort;
}
do { do {
arrange_input_buffer(&self->zst, &ibuflen); arrange_input_buffer(&self->zst, &ibuflen);
flush = ibuflen == 0 ? Z_FINISH : Z_NO_FLUSH; flush = ibuflen == 0 ? Z_FINISH : Z_NO_FLUSH;
do { do {
length = arrange_output_buffer(&self->zst, &RetVal, length); if (self->zst.avail_out == 0) {
if (length < 0) if (Buffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0)
goto abort; goto abort;
}
Py_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS
err = inflate(&self->zst, flush); err = inflate(&self->zst, flush);
@ -1193,13 +1243,14 @@ zlib_Decompress_flush_impl(compobject *self, PyTypeObject *cls,
} }
} }
if (_PyBytes_Resize(&RetVal, self->zst.next_out - RetVal = Buffer_Finish(&buffer, self->zst.avail_out);
(Byte *)PyBytes_AS_STRING(RetVal)) == 0) { if (RetVal != NULL) {
goto success; goto success;
} }
abort: abort:
Py_CLEAR(RetVal); Buffer_OnError(&buffer);
RetVal = NULL;
success: success:
PyBuffer_Release(&data); PyBuffer_Release(&data);
LEAVE_ZLIB(self); LEAVE_ZLIB(self);