BZ2Decompressor and zlib._ZlibDecompressor set needs_input=True while output is still buffered #159999

Bug description

When max_length cuts the output short and the call also consumes the whole input buffer, bz2.BZ2Decompressor and zlib._ZlibDecompressor set needs_input to True even though a decompress(b'') call would return more data. The documented meaning of the flag is the opposite: False means "there is buffered output, call again with b''"; True means "all output has been produced, feed more input".

>>> import bz2, zlib
>>> stream = bz2.compress(b"\x00\x10 0@" * 2)
>>> d = bz2.BZ2Decompressor()
>>> d.decompress(stream[:40], max_length=0)
b''
>>> d.needs_input, d.eof
(True, False)
>>> d.decompress(b"")   # Expecting: b'' if needs_input is to be believed
b'\x00\x10 0@\x00\x10 0@'

>>> d = zlib._ZlibDecompressor()
>>> d.decompress(zlib.compress(b"\x00")[:4], max_length=0)
b''
>>> d.needs_input
True
>>> d.decompress(b"")
b'\x00'

lzma.LZMADecompressor and compression.zstd.ZstdDecompressor get this right. Compare lzma on the same kind of input:

>>> import lzma
>>> d = lzma.LZMADecompressor()
>>> d.decompress(lzma.compress(b"\x00" * 21), max_length=2)
b'\x00\x00'
>>> d.needs_input
False

The cause is that both modules decide needs_input from the input buffer alone. In _bz2module.c the branch that handles "input exhausted" unconditionally stores 1: Modules/_bz2module.c lines 539-542. zlibmodule.c has the same branch for _ZlibDecompressor: Modules/zlibmodule.c lines 1646-1649. _lzmamodule.c checks whether the output buffer was also full before it decides, with a comment explaining why: Modules/_lzmamodule.c lines 1069-1082. Porting that check to the other two modules looks like the whole fix.

As a practical matter I don't see a way for this to lose data for a caller following the documented protocol: the stream's trailer has not been consumed yet in these states, so the caller still has input to feed and will get the buffered output then. It does make needs_input unreliable as a "no more buffered output" signal, which is what max_length users (including DecompressReader in compression/_common/_streams.py, which raises EOFError rather than call decompress(b'') when needs_input is true at end of file) rely on. So it is worth fixing, but it is not urgent.

Found by the Hypothesis tests in #159960, which record it as test_needs_input_means_drained under expectedFailure for the bz2 and zlib backends. The reproducers depend on where the compressor places block boundaries; the ones above are for zlib 1.3.2 and libbz2 1.0.8.

CPython versions tested on

CPython main (4e3d5ec)

Operating systems tested on

Linux