Skip to content

add set_stream to DecompIO #11183

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
Open

add set_stream to DecompIO #11183

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7dc82cf
added a reset function to allow buffer reuse. see:
keredson Apr 2, 2023
81d82fb
verify old and new buffers are the same size
keredson Apr 2, 2023
489d121
changed the new function name to set_stream() (reset is too generic)
keredson Apr 2, 2023
6c4b630
fixed doc formatting error
keredson Apr 2, 2023
3676573
code formatter
keredson Apr 2, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions docs/library/zlib.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,3 +36,15 @@ Functions

This class is MicroPython extension. It's included on provisional
basis and may be changed considerably or removed in later versions.

.. function:: DecompIO.set_stream(obj, stream, wbits=0, /)

Update a `DecompIO` to use a new stream object. This is exactly like
creating a new `DecompIO` object, except the buffer (usually 32k) is
reused and therefore doesn't have to be reallocated. This can
prevent a random `MemoryError` due to memory fragmentation when
processing large numbers of streams.

The buffer sizes of the two streams must match, or this will throw
a `ValueError`.

70 changes: 56 additions & 14 deletions extmod/moduzlib.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,21 +66,8 @@ STATIC int read_src_stream(TINF_DATA *data) {
return c;
}

STATIC mp_obj_t decompio_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
mp_arg_check_num(n_args, n_kw, 1, 2, false);
mp_get_stream_raise(args[0], MP_STREAM_OP_READ);
mp_obj_decompio_t *o = mp_obj_malloc(mp_obj_decompio_t, type);
memset(&o->decomp, 0, sizeof(o->decomp));
o->decomp.readSource = read_src_stream;
o->src_stream = args[0];
o->eof = false;

mp_int_t dict_opt = 0;
STATIC uint calc_dict_sz(mp_int_t dict_opt, mp_obj_decompio_t *o) {
uint dict_sz;
if (n_args > 1) {
dict_opt = mp_obj_get_int(args[1]);
}

if (dict_opt >= 16) {
int st = uzlib_gzip_parse_header(&o->decomp);
if (st != TINF_OK) {
Expand All @@ -100,6 +87,23 @@ STATIC mp_obj_t decompio_make_new(const mp_obj_type_t *type, size_t n_args, size
} else {
dict_sz = 1 << -dict_opt;
}
return dict_sz;
}

STATIC mp_obj_t decompio_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
mp_arg_check_num(n_args, n_kw, 1, 2, false);
mp_get_stream_raise(args[0], MP_STREAM_OP_READ);
mp_obj_decompio_t *o = mp_obj_malloc(mp_obj_decompio_t, type);
memset(&o->decomp, 0, sizeof(o->decomp));
o->decomp.readSource = read_src_stream;
o->src_stream = args[0];
o->eof = false;

mp_int_t dict_opt = 0;
if (n_args > 1) {
dict_opt = mp_obj_get_int(args[1]);
}
uint dict_sz = calc_dict_sz(dict_opt, o);

uzlib_uncompress_init(&o->decomp, m_new(byte, dict_sz), dict_sz);
return MP_OBJ_FROM_PTR(o);
Expand All @@ -125,9 +129,47 @@ STATIC mp_uint_t decompio_read(mp_obj_t o_in, void *buf, mp_uint_t size, int *er
return o->decomp.dest - (byte *)buf;
}

STATIC mp_obj_t mod_uzlib_set_stream(size_t n_args, const mp_obj_t *args) {
mp_get_stream_raise(args[1], MP_STREAM_OP_READ);
mp_obj_decompio_t *o = args[0];
TINF_DATA *decomp = &o->decomp;
unsigned char *dict_ring = decomp->dict_ring;
unsigned int dict_size = decomp->dict_size;
memset(&o->decomp, 0, sizeof(o->decomp));
o->decomp.readSource = read_src_stream;
o->src_stream = args[1];
o->eof = false;
decomp->dict_ring = dict_ring;
decomp->dict_size = dict_size;

mp_int_t dict_opt = 0;
if (n_args > 2) {
dict_opt = mp_obj_get_int(args[2]);
}
uint dict_sz = calc_dict_sz(dict_opt, o);

if (dict_sz != dict_size) {
mp_raise_ValueError(MP_ERROR_TEXT("compression header buffer sizes must match (to reuse buffer)"));
}

// dict_opt = uzlib_zlib_parse_header(decomp);
for (uint i = 0; i < dict_size; ++i) {
dict_ring[i] = 0;
}
decomp->eof = 0;
decomp->bitcount = 0;
decomp->bfinal = 0;
decomp->btype = -1;
decomp->dict_idx = 0;
decomp->curlen = 0;
return o;
}
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_uzlib_set_stream_obj, 1, 3, mod_uzlib_set_stream);

#if !MICROPY_ENABLE_DYNRUNTIME
STATIC const mp_rom_map_elem_t decompio_locals_dict_table[] = {
{ MP_ROM_QSTR(MP_QSTR_read), MP_ROM_PTR(&mp_stream_read_obj) },
{ MP_ROM_QSTR(MP_QSTR_set_stream), MP_ROM_PTR(&mod_uzlib_set_stream_obj) },
{ MP_ROM_QSTR(MP_QSTR_readinto), MP_ROM_PTR(&mp_stream_readinto_obj) },
{ MP_ROM_QSTR(MP_QSTR_readline), MP_ROM_PTR(&mp_stream_unbuffered_readline_obj) },
};
Expand Down
5 changes: 5 additions & 0 deletions tests/extmod/uzlib_decompio.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,3 +31,8 @@
print(inp.read())
except OSError as e:
print(repr(e))

inp = zlib.DecompIO(io.BytesIO(b"x\x9c+.)JM\xcc5\x04\x00\x0b\xe0\x02\xbe"))
assert inp.read() == b"stream1"
inp.set_stream(io.BytesIO(b"x\x9c+.)JM\xcc5\x02\x00\x0b\xe1\x02\xbf"))
assert inp.read() == b"stream2"