diff --git a/Cargo.lock b/Cargo.lock index d51970db471..9600a51ba68 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3326,6 +3326,7 @@ dependencies = [ "ascii", "base64", "bitflags 2.13.1", + "bzip2", "crc32fast", "getrandom 0.4.3", "itertools 0.15.0", @@ -3571,7 +3572,6 @@ dependencies = [ "adler32", "ascii", "blake2", - "bzip2", "constant_time_eq", "crossbeam-utils", "csv-core", diff --git a/Lib/test/test_bz2.py b/Lib/test/test_bz2.py index bcd4e033b59..64293d75733 100644 --- a/Lib/test/test_bz2.py +++ b/Lib/test/test_bz2.py @@ -1032,7 +1032,6 @@ def test_failure(self): # Previously, a second call could crash due to internal inconsistency self.assertRaises(Exception, bzd.decompress, self.BAD_DATA * 30) - @unittest.expectedFailure # TODO: RUSTPYTHON; Wrong error message def test_decompress_after_data_error(self): data = bytes.fromhex( "425a6839314159265359000000000000007fffff000000000000000000000000" diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index 21b700a5526..72d7ce4a6f7 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -11,6 +11,7 @@ license.workspace = true [features] default = ["std"] binascii = ["dep:base64", "dep:crc32fast"] +bz2 = ["std", "dep:bzip2"] cjk-codecs = [] inet = ["std"] json = ["dep:memchr", "std"] @@ -28,6 +29,7 @@ rustpython-wtf8 = { workspace = true } ascii = { workspace = true } base64 = { workspace = true, optional = true } bitflags = { workspace = true } +bzip2 = { workspace = true, optional = true } crc32fast = { workspace = true, optional = true } getrandom = { workspace = true } itertools = { workspace = true } diff --git a/crates/common/src/compression/bz2.rs b/crates/common/src/compression/bz2.rs new file mode 100644 index 00000000000..d346ce4bd9c --- /dev/null +++ b/crates/common/src/compression/bz2.rs @@ -0,0 +1,324 @@ +// spell-checker:ignore chunker libbz +//! VM-independent bzip2 stream engine. +//! +//! The engine owns its native stream state and reports plain Rust errors so +//! interpreter and embedding layers can provide their own object and exception +//! adapters. + +use bzip2::{Action, Compress, Compression, Decompress, Error, Status}; + +use super::Chunker; + +const INITIAL_BUFFER_SIZE: usize = 8192; +const BIGCHUNK: usize = 512 * 1024; + +/// Double the output block until `BIGCHUNK`, then keep the size fixed. +const fn new_buffer_size(current_size: usize) -> usize { + if current_size < BIGCHUNK { + current_size + current_size + } else { + current_size + } +} + +/// Stream errors reported by the bzip2 engine. Exception mapping belongs to +/// the interpreter adapter. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum Bz2Error { + /// `BZ_PARAM_ERROR` + Param, + /// `BZ_DATA_ERROR` / `BZ_DATA_ERROR_MAGIC` + Data, + /// `BZ_SEQUENCE_ERROR` + Sequence, + /// `BZ_MEM_ERROR` + Mem, +} + +impl From for Bz2Error { + fn from(error: Error) -> Self { + match error { + Error::Param => Self::Param, + Error::Data | Error::DataMagic => Self::Data, + Error::Sequence => Self::Sequence, + } + } +} + +/// Incremental compressor. After `flush`, the stream is finished and must +/// not be used again. +pub struct Compressor { + compress: Compress, + flushed: bool, +} + +impl Compressor { + /// Create a compressor. Returns `None` when `compresslevel` is outside + /// `1..=9`. A zero work factor selects libbz2's default of 30. + pub fn new(compresslevel: i64) -> Option { + let level = u32::try_from(compresslevel) + .ok() + .and_then(Compression::try_new)?; + Some(Self { + compress: Compress::new(level, 0), + flushed: false, + }) + } + + #[must_use] + pub fn is_flushed(&self) -> bool { + self.flushed + } + + /// Compress `data` without finishing the stream. + pub fn compress(&mut self, data: &[u8]) -> Result, Bz2Error> { + self.run(data, Action::Run) + } + + /// Finish the stream. The compressor must not be used again. + pub fn flush(&mut self) -> Result, Bz2Error> { + self.flushed = true; + self.run(&[], Action::Finish) + } + + /// One pass over the input with the requested action, growing the output + /// block whenever libbz2 fills it. + fn run(&mut self, mut input: &[u8], action: Action) -> Result, Bz2Error> { + let mut out = Vec::new(); + let mut block = vec![0u8; INITIAL_BUFFER_SIZE]; + loop { + // In regular compression mode, stop when input data is exhausted. + if action == Action::Run && input.is_empty() { + break; + } + let previous_in = self.compress.total_in(); + let previous_out = self.compress.total_out(); + let status = self.compress.compress(input, &mut block, action)?; + let consumed = (self.compress.total_in() - previous_in) as usize; + let produced = (self.compress.total_out() - previous_out) as usize; + out.extend_from_slice(&block[..produced]); + input = &input[consumed..]; + // In flushing mode, stop when all buffered data has been flushed. + if action == Action::Finish && status == Status::StreamEnd { + break; + } + if produced == block.len() { + block = vec![0u8; new_buffer_size(block.len())]; + } + } + out.shrink_to_fit(); + Ok(out) + } +} + +/// Incremental decompressor. The first failure is latched so later calls +/// can be refused by the caller; re-entering the native stream after a +/// failure can write out of bounds. +pub struct Decompressor { + decompress: Decompress, + eof: bool, + failed: bool, + needs_input: bool, + unused_data: Vec, + /// Input handed to `decompress` that libbz2 has not consumed yet. + input_buffer: Vec, +} + +impl Decompressor { + /// Create a decompressor using the fast (non-`small`) algorithm. + #[must_use] + pub fn new() -> Self { + Self { + decompress: Decompress::new(false), + eof: false, + failed: false, + needs_input: true, + unused_data: Vec::new(), + input_buffer: Vec::new(), + } + } + + #[must_use] + pub fn eof(&self) -> bool { + self.eof + } + + #[must_use] + pub fn failed(&self) -> bool { + self.failed + } + + #[must_use] + pub fn needs_input(&self) -> bool { + self.needs_input + } + + #[must_use] + pub fn unused_data(&self) -> &[u8] { + &self.unused_data + } + + /// Decompress from new input and any bytes buffered by an earlier call. + /// `max_length` of `None` means unlimited output. + pub fn decompress( + &mut self, + data: &[u8], + max_length: Option, + ) -> Result, Bz2Error> { + let max_length = max_length.unwrap_or(usize::MAX); + let mut out = Vec::new(); + let mut block = vec![0u8; INITIAL_BUFFER_SIZE.min(max_length)]; + + let mut failed = None; + let mut stream_end = false; + let leftover = { + let mut chunks = Chunker::chain(&self.input_buffer, data); + loop { + let chunk = chunks.chunk(); + let previous_in = self.decompress.total_in(); + let previous_out = self.decompress.total_out(); + let status = self.decompress.decompress(chunk, &mut block); + let consumed = (self.decompress.total_in() - previous_in) as usize; + let produced = (self.decompress.total_out() - previous_out) as usize; + chunks.advance(consumed); + out.extend_from_slice(&block[..produced]); + match status { + Err(error) => { + failed = Some(error.into()); + break; + } + Ok(Status::MemNeeded) => { + failed = Some(Bz2Error::Mem); + break; + } + Ok(Status::StreamEnd) => { + stream_end = true; + break; + } + Ok(_) => {} + } + if chunks.is_empty() { + break; + } + if produced == block.len() { + // The output block is full: grow it unless `max_length` + // has already been reached. + if out.len() == max_length { + break; + } + block = vec![0u8; new_buffer_size(block.len()).min(max_length - out.len())]; + } + } + if chunks.is_empty() { + None + } else { + Some(chunks.to_vec()) + } + }; + + if let Some(error) = failed { + return Err(self.fail(error)); + } + + if stream_end { + self.eof = true; + self.needs_input = false; + self.input_buffer.clear(); + if let Some(unused) = leftover { + self.unused_data = unused; + } + } else if let Some(remaining) = leftover { + self.needs_input = false; + self.input_buffer = remaining; + } else { + self.needs_input = true; + self.input_buffer.clear(); + } + out.shrink_to_fit(); + Ok(out) + } + + /// Latch the first failure and drop the pending input with it. + fn fail(&mut self, error: Bz2Error) -> Bz2Error { + self.failed = true; + self.needs_input = false; + self.input_buffer = Vec::new(); + error + } +} + +impl Default for Decompressor { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn roundtrip(data: &[u8], level: i64) -> Vec { + let mut compressor = Compressor::new(level).unwrap(); + let mut encoded = compressor.compress(data).unwrap(); + encoded.extend(compressor.flush().unwrap()); + let mut decompressor = Decompressor::new(); + let out = decompressor.decompress(&encoded, None).unwrap(); + assert!(decompressor.eof()); + out + } + + #[test] + fn invalid_level_is_rejected() { + assert!(Compressor::new(0).is_none()); + assert!(Compressor::new(10).is_none()); + assert!(Compressor::new(-1).is_none()); + } + + #[test] + fn streaming_roundtrip() { + let data = b"the quick brown fox jumps over the lazy dog".repeat(50); + assert_eq!(roundtrip(&data, 9), data); + } + + #[test] + fn unused_data_after_stream_end() { + let mut compressor = Compressor::new(9).unwrap(); + let mut encoded = compressor.compress(b"hello").unwrap(); + encoded.extend(compressor.flush().unwrap()); + encoded.extend_from_slice(b"trailing"); + + let mut decompressor = Decompressor::new(); + let out = decompressor.decompress(&encoded, None).unwrap(); + assert_eq!(out, b"hello"); + assert!(decompressor.eof()); + assert!(!decompressor.needs_input()); + assert_eq!(decompressor.unused_data(), b"trailing"); + } + + #[test] + fn max_length_leaves_unconsumed_input() { + let data = b"abcdefghij".repeat(20); + let mut compressor = Compressor::new(9).unwrap(); + let mut encoded = compressor.compress(&data).unwrap(); + encoded.extend(compressor.flush().unwrap()); + + let mut decompressor = Decompressor::new(); + let first = decompressor.decompress(&encoded, Some(5)).unwrap(); + assert_eq!(first.len(), 5); + assert!(!decompressor.eof()); + assert!(!decompressor.needs_input()); + let rest = decompressor.decompress(&[], None).unwrap(); + assert_eq!([first, rest].concat(), data); + assert!(decompressor.eof()); + } + + #[test] + fn bad_data_latches_failure() { + let mut decompressor = Decompressor::new(); + let err = decompressor.decompress(b"not a bz2 stream", None); + assert_eq!(err, Err(Bz2Error::Data)); + assert!(decompressor.failed()); + assert!(!decompressor.needs_input()); + assert!(!decompressor.eof()); + } +} diff --git a/crates/common/src/compression/mod.rs b/crates/common/src/compression/mod.rs index f0d2b7b5f19..9f84d2590dd 100644 --- a/crates/common/src/compression/mod.rs +++ b/crates/common/src/compression/mod.rs @@ -69,6 +69,8 @@ impl<'a> Chunker<'a> { } } +#[cfg(feature = "bz2")] +pub mod bz2; #[cfg(all( feature = "lzma", not(any(target_os = "android", target_arch = "wasm32")) diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 575bf64e365..ddf0cf91558 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -11,7 +11,7 @@ pub mod binascii; pub mod borrow; pub mod boxvec; pub mod cformat; -#[cfg(any(feature = "lzma", feature = "zlib"))] +#[cfg(any(feature = "bz2", feature = "lzma", feature = "zlib"))] pub mod compression; pub mod encodings; pub mod float_ops; diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index cd955597d94..794b8eee176 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -30,7 +30,7 @@ __ssl-rustls = ["ssl", "rustls", "rustls-native-certs", "rustls-pemfile", "rustl # rustpython crates rustpython-derive = { workspace = true } rustpython-vm = { workspace = true, default-features = false, features = ["compiler"]} -rustpython-common = { workspace = true, features = ["binascii", "cjk-codecs", "inet", "json", "lzma", "zlib"] } +rustpython-common = { workspace = true, features = ["binascii", "bz2", "cjk-codecs", "inet", "json", "lzma", "zlib"] } rustpython-host_env = { workspace = true } rustpython-unicode = { workspace = true } @@ -79,7 +79,6 @@ constant_time_eq = { workspace = true } # compression adler32 = { workspace = true } -bzip2 = { workspace = true } # tkinter jiff = { workspace = true } diff --git a/crates/stdlib/src/bz2.rs b/crates/stdlib/src/bz2.rs index fcde54057e5..25efa228d60 100644 --- a/crates/stdlib/src/bz2.rs +++ b/crates/stdlib/src/bz2.rs @@ -4,54 +4,50 @@ pub(crate) use _bz2::module_def; #[pymodule] mod _bz2 { - use crate::compression::{ - DecompressArgs, DecompressError, DecompressState, DecompressStatus, Decompressor, - }; + use crate::compression::DecompressArgs; use crate::vm::{ Py, VirtualMachine, - builtins::{PyBytesRef, PyType}, + builtins::{PyBaseExceptionRef, PyBytesRef, PyType}, common::lock::PyMutex, function::{ArgBytesLike, OptionalArg}, object::PyResult, types::Constructor, }; use alloc::fmt; - use bzip2::{Decompress, Status, write::BzEncoder}; - use core::mem; - use rustpython_vm::convert::ToPyException; - use std::io::Write; - - const BUFSIZ: usize = 8192; + use rustpython_common::compression::bz2 as backend; - #[pyattr] - #[pyclass(name = "BZ2Decompressor")] - #[derive(PyPayload)] - struct BZ2Decompressor { - state: PyMutex>, + fn map_bz2_error(error: backend::Bz2Error, vm: &VirtualMachine) -> PyBaseExceptionRef { + match error { + backend::Bz2Error::Param => { + vm.new_value_error("Internal error - invalid parameters passed to libbzip2") + } + backend::Bz2Error::Data => vm.new_os_error("Invalid data stream"), + backend::Bz2Error::Sequence => vm.new_runtime_error( + "Internal error - Invalid sequence of commands sent to libbzip2", + ), + backend::Bz2Error::Mem => vm.new_memory_error("out of memory"), + } } - impl Decompressor for Decompress { - type Flush = (); - type Status = Status; - type Error = bzip2::Error; + struct PyBZ2DecompressorInner { + decompress: backend::Decompressor, + unused_data: PyBytesRef, + } - fn total_in(&self) -> u64 { - self.total_in() - } - fn decompress_vec( - &mut self, - input: &[u8], - output: &mut Vec, - (): Self::Flush, - ) -> Result { - self.decompress_vec(input, output) + impl PyBZ2DecompressorInner { + fn sync_visible_state(&mut self, vm: &VirtualMachine) { + if self.unused_data.as_bytes() != self.decompress.unused_data() { + self.unused_data = vm.ctx.new_bytes(self.decompress.unused_data().to_vec()); + } } } - impl DecompressStatus for Status { - fn is_stream_end(&self) -> bool { - *self == Self::StreamEnd - } + #[pyattr] + #[pyclass(name = "BZ2Decompressor", traverse)] + #[derive(PyPayload)] + struct BZ2Decompressor { + #[pytraverse(skip)] + inner: PyMutex, } impl fmt::Debug for BZ2Decompressor { @@ -65,7 +61,10 @@ mod _bz2 { fn py_new(_cls: &Py, _args: Self::Args, vm: &VirtualMachine) -> PyResult { Ok(Self { - state: PyMutex::new(DecompressState::new(Decompress::new(false), vm)), + inner: PyMutex::new(PyBZ2DecompressorInner { + decompress: backend::Decompressor::new(), + unused_data: vm.ctx.empty_bytes.clone(), + }), }) } } @@ -77,50 +76,44 @@ mod _bz2 { let max_length = args.max_length(); let data = &*args.data(); - let mut state = self.state.lock(); - state - .decompress(data, max_length, BUFSIZ, vm) - .map_err(|e| match e { - DecompressError::Decompress(err) => vm.new_os_error(err.to_string()), - DecompressError::Eof(err) => err.to_pyexception(vm), - }) + let mut inner = self.inner.lock(); + if inner.decompress.eof() { + return Err(vm.new_eof_error("End of stream already reached")); + } + if inner.decompress.failed() { + return Err(vm.new_value_error("Decompressor is unusable after a previous error")); + } + let result = inner.decompress.decompress(data, max_length); + inner.sync_visible_state(vm); + result.map_err(|error| map_bz2_error(error, vm)) } #[pygetset] fn eof(&self) -> bool { - self.state.lock().eof() + self.inner.lock().decompress.eof() } #[pygetset] fn unused_data(&self) -> PyBytesRef { - self.state.lock().unused_data() + self.inner.lock().unused_data.clone() } #[pygetset] fn needs_input(&self) -> bool { - // False if the decompress() method can provide more - // decompressed data before requiring new uncompressed input. - self.state.lock().needs_input() + self.inner.lock().decompress.needs_input() } #[pymethod(name = "__reduce__")] fn reduce(&self, vm: &VirtualMachine) -> PyResult<()> { Err(vm.new_type_error("cannot pickle '_bz2.BZ2Decompressor' object")) } - - // TODO: mro()? - } - - struct CompressorState { - flushed: bool, - encoder: Option>>, } #[pyattr] #[pyclass(name = "BZ2Compressor")] #[derive(PyPayload)] struct BZ2Compressor { - state: PyMutex, + state: PyMutex, } impl fmt::Debug for BZ2Compressor { @@ -137,21 +130,11 @@ mod _bz2 { (compresslevel,): Self::Args, vm: &VirtualMachine, ) -> PyResult { - // TODO: seriously? - // compresslevel.unwrap_or(bzip2::Compression::best().level().try_into().unwrap()); let compresslevel = compresslevel.unwrap_or(9); - let level = match compresslevel { - valid_level @ 1..=9 => bzip2::Compression::new(valid_level as u32), - _ => { - return Err(vm.new_value_error("compresslevel must be between 1 and 9")); - } - }; - + let compressor = backend::Compressor::new(i64::from(compresslevel)) + .ok_or_else(|| vm.new_value_error("compresslevel must be between 1 and 9"))?; Ok(Self { - state: PyMutex::new(CompressorState { - flushed: false, - encoder: Some(BzEncoder::new(Vec::new(), level)), - }), + state: PyMutex::new(compressor), }) } } @@ -159,34 +142,22 @@ mod _bz2 { #[pyclass(with(Constructor))] impl BZ2Compressor { #[pymethod] - fn compress(&self, data: ArgBytesLike, vm: &VirtualMachine) -> PyResult { - let mut state = self.state.lock(); - if state.flushed { + fn compress(&self, data: ArgBytesLike, vm: &VirtualMachine) -> PyResult> { + let mut compressor = self.state.lock(); + if compressor.is_flushed() { return Err(vm.new_value_error("Compressor has been flushed")); } - - let CompressorState { encoder, .. } = &mut *state; - let encoder = encoder.as_mut().unwrap(); - data.with_ref(|input_bytes| encoder.write_all(input_bytes).unwrap()); - // BzEncoder writes its pending output at the start of the next write. - assert_eq!(encoder.write(&[]).unwrap(), 0); - Ok(vm.ctx.new_bytes(mem::take(encoder.get_mut()))) + data.with_ref(|input| compressor.compress(input)) + .map_err(|error| map_bz2_error(error, vm)) } #[pymethod] - fn flush(&self, vm: &VirtualMachine) -> PyResult { - let mut state = self.state.lock(); - if state.flushed { + fn flush(&self, vm: &VirtualMachine) -> PyResult> { + let mut compressor = self.state.lock(); + if compressor.is_flushed() { return Err(vm.new_value_error("Repeated call to flush()")); } - - // let CompressorState { flushed, encoder } = &mut *state; - let CompressorState { encoder, .. } = &mut *state; - - // TODO: handle Err - let out = encoder.take().unwrap().finish().unwrap(); - state.flushed = true; - Ok(vm.ctx.new_bytes(out.to_vec())) + compressor.flush().map_err(|error| map_bz2_error(error, vm)) } #[pymethod(name = "__reduce__")] diff --git a/crates/stdlib/src/compression.rs b/crates/stdlib/src/compression.rs index 563a32f7b35..3ca8c6d10eb 100644 --- a/crates/stdlib/src/compression.rs +++ b/crates/stdlib/src/compression.rs @@ -1,14 +1,6 @@ -// spell-checker:ignore chunker - //! internal shared module for compression libraries use crate::vm::function::{ArgBytesLike, ArgSize, OptionalArg}; -use crate::vm::{ - VirtualMachine, - builtins::{PyBaseExceptionRef, PyBytesRef}, - convert::ToPyException, -}; -use rustpython_common::compression::Chunker; #[derive(FromArgs)] pub(crate) struct DecompressArgs { @@ -33,189 +25,3 @@ impl DecompressArgs { .and_then(|ArgSize { value }| usize::try_from(value).ok()) } } - -pub(crate) trait Decompressor { - type Flush: DecompressFlushKind; - type Status: DecompressStatus; - type Error; - - fn total_in(&self) -> u64; - fn decompress_vec( - &mut self, - input: &[u8], - output: &mut Vec, - flush: Self::Flush, - ) -> Result; - fn maybe_set_dict(&mut self, err: Self::Error) -> Result<(), Self::Error> { - Err(err) - } -} - -pub(crate) trait DecompressStatus { - fn is_stream_end(&self) -> bool; -} - -pub(crate) trait DecompressFlushKind: Copy { - const SYNC: Self; -} - -impl DecompressFlushKind for () { - const SYNC: Self = (); -} - -pub(crate) const fn flush_sync(_final_chunk: bool) -> T { - T::SYNC -} - -pub(crate) fn _decompress_chunks( - data: &mut Chunker<'_>, - d: &mut D, - bufsize: usize, - max_length: Option, - calc_flush: impl Fn(bool) -> D::Flush, -) -> Result<(Vec, bool), D::Error> { - if data.is_empty() { - return Ok((Vec::new(), true)); - } - let max_length = max_length.unwrap_or(usize::MAX); - let mut buf = Vec::new(); - - 'outer: loop { - let chunk = data.chunk(); - let flush = calc_flush(chunk.len() == data.len()); - loop { - let additional = core::cmp::min(bufsize, max_length - buf.capacity()); - if additional == 0 { - return Ok((buf, false)); - } - buf.reserve_exact(additional); - - let prev_in = d.total_in(); - let res = d.decompress_vec(chunk, &mut buf, flush); - let consumed = d.total_in() - prev_in; - - data.advance(consumed as usize); - - match res { - Ok(status) => { - let stream_end = status.is_stream_end(); - if stream_end || data.is_empty() { - // we've reached the end of the stream, we're done - buf.shrink_to_fit(); - return Ok((buf, stream_end)); - } else if !chunk.is_empty() && consumed == 0 { - // we're gonna need a bigger buffer - continue; - } - // next chunk - continue 'outer; - } - Err(e) => { - d.maybe_set_dict(e)?; - // now try the next chunk - continue 'outer; - } - }; - } - } -} - -#[derive(Debug)] -pub(crate) struct DecompressState { - decompress: D, - unused_data: PyBytesRef, - input_buffer: Vec, - eof: bool, - needs_input: bool, -} - -impl DecompressState { - pub(crate) fn new(decompress: D, vm: &VirtualMachine) -> Self { - Self { - decompress, - unused_data: vm.ctx.empty_bytes.clone(), - input_buffer: Vec::new(), - eof: false, - needs_input: true, - } - } - - pub(crate) const fn eof(&self) -> bool { - self.eof - } - - pub(crate) fn unused_data(&self) -> PyBytesRef { - self.unused_data.clone() - } - - pub(crate) const fn needs_input(&self) -> bool { - self.needs_input - } - - pub(crate) fn decompress( - &mut self, - data: &[u8], - max_length: Option, - bufsize: usize, - vm: &VirtualMachine, - ) -> Result, DecompressError> { - if self.eof { - return Err(DecompressError::Eof(EofError)); - } - - let input_buffer = &mut self.input_buffer; - let d = &mut self.decompress; - - let mut chunks = Chunker::chain(input_buffer, data); - - let prev_len = chunks.len(); - let (ret, stream_end) = - match _decompress_chunks(&mut chunks, d, bufsize, max_length, flush_sync) { - Ok((buf, stream_end)) => (Ok(buf), stream_end), - Err(err) => (Err(err), false), - }; - let consumed = prev_len - chunks.len(); - - self.eof |= stream_end; - - if self.eof { - self.needs_input = false; - if !chunks.is_empty() { - self.unused_data = vm.ctx.new_bytes(chunks.to_vec()); - } - } else if chunks.is_empty() { - input_buffer.clear(); - self.needs_input = true; - } else { - self.needs_input = false; - if let Some(n_consumed_from_data) = consumed.checked_sub(input_buffer.len()) { - input_buffer.clear(); - input_buffer.extend_from_slice(&data[n_consumed_from_data..]); - } else { - input_buffer.drain(..consumed); - input_buffer.extend_from_slice(data); - } - } - - ret.map_err(DecompressError::Decompress) - } -} - -pub(crate) enum DecompressError { - Decompress(E), - Eof(EofError), -} - -impl From for DecompressError { - fn from(err: E) -> Self { - Self::Decompress(err) - } -} - -pub(crate) struct EofError; - -impl ToPyException for EofError { - fn to_pyexception(&self, vm: &VirtualMachine) -> PyBaseExceptionRef { - vm.new_eof_error("End of stream already reached") - } -}