diff --git a/Cargo.lock b/Cargo.lock index 1137c15b9..c5fe5809c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -150,6 +150,27 @@ version = "3.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f1e260c3a9040a7c19a12468758f4c16f31a81a1fe087482be9570ec864bb6c" +[[package]] +name = "bzip2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6afcd980b5f3a45017c57e57a2fcccbb351cc43a356ce117ef760ef8052b89b0" +dependencies = [ + "bzip2-sys", + "libc", +] + +[[package]] +name = "bzip2-sys" +version = "0.1.11+1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "736a955f3fa7875102d57c82b8cac37ec45224a07fd32d58f9f7a186b6cd4cdc" +dependencies = [ + "cc", + "libc", + "pkg-config", +] + [[package]] name = "caseless" version = "0.2.1" @@ -1722,6 +1743,7 @@ dependencies = [ "ascii", "base64", "blake2", + "bzip2", "cfg-if", "crc32fast", "crossbeam-utils", diff --git a/stdlib/Cargo.toml b/stdlib/Cargo.toml index 90bab2a8c..91c3cd764 100644 --- a/stdlib/Cargo.toml +++ b/stdlib/Cargo.toml @@ -41,6 +41,7 @@ unic-ucd-ident = "0.9.0" adler32 = "1.2.0" crc32fast = "1.3.2" flate2 = "1.0.23" +bzip2 = "0.4" num-complex = "0.4.0" num-bigint = "0.4.3" diff --git a/stdlib/src/bz2.rs b/stdlib/src/bz2.rs new file mode 100644 index 000000000..977dea918 --- /dev/null +++ b/stdlib/src/bz2.rs @@ -0,0 +1,246 @@ +pub(crate) use _bz2::make_module; + +#[pymodule] +mod _bz2 { + use crate::common::lock::PyMutex; + use crate::vm::{ + builtins::{PyBytesRef, PyTypeRef}, + function::{ArgBytesLike, OptionalArg}, + object::{PyPayload, PyResult}, + types::Constructor, + VirtualMachine, + }; + use bzip2::{write::BzEncoder, Decompress, Status}; + use std::{fmt, io::Write}; + + // const BUFSIZ: i32 = 8192; + + struct DecompressorState { + decoder: Decompress, + eof: bool, + needs_input: bool, + // input_buffer: Vec, + // output_buffer: Vec, + } + + #[pyattr] + #[pyclass(name = "BZ2Decompressor")] + #[derive(PyPayload)] + struct BZ2Decompressor { + state: PyMutex, + } + + impl fmt::Debug for BZ2Decompressor { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "_bz2.BZ2Decompressor") + } + } + + impl Constructor for BZ2Decompressor { + type Args = (); + + fn py_new(cls: PyTypeRef, _: Self::Args, vm: &VirtualMachine) -> PyResult { + Self { + state: PyMutex::new(DecompressorState { + decoder: Decompress::new(false), + eof: false, + needs_input: true, + // input_buffer: Vec::new(), + // output_buffer: Vec::new(), + }), + } + .into_ref_with_type(vm, cls) + .map(Into::into) + } + } + + #[pyclass(with(Constructor))] + impl BZ2Decompressor { + #[pymethod] + fn decompress( + &self, + data: ArgBytesLike, + // TODO: PyIntRef + max_length: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + let max_length = max_length.unwrap_or(-1); + if max_length >= 0 { + return Err(vm.new_not_implemented_error( + "the max_value argument is not implemented yet".to_owned(), + )); + } + // let max_length = if max_length < 0 || max_length >= BUFSIZ { + // BUFSIZ + // } else { + // max_length + // }; + + let mut state = self.state.lock(); + let DecompressorState { + decoder, + eof, + .. + // needs_input, + // input_buffer, + // output_buffer, + } = &mut *state; + + if *eof { + return Err(vm.new_exception_msg( + vm.ctx.exceptions.eof_error.to_owned(), + "End of stream already reached".to_owned(), + )); + } + + // data.with_ref(|data| input_buffer.extend(data)); + + // If max_length is negative: + // read the input X bytes at a time, compress it and append it to output. + // Once you're out of input, setting needs_input to true and return the + // output as bytes. + // + // TODO: + // If max_length is non-negative: + // Read the input X bytes at a time, compress it and append it to + // the output. If output reaches `max_length` in size, return + // it (up to max_length), and store the rest of the output + // for later. + + // TODO: arbitrary choice, not the right way to do it. + let mut buf = Vec::with_capacity(data.len() * 32); + + let before = decoder.total_in(); + let res = data.with_ref(|data| decoder.decompress_vec(data, &mut buf)); + let _written = (decoder.total_in() - before) as usize; + + let res = match res { + Ok(x) => x, + // TODO: error message + _ => return Err(vm.new_os_error("Invalid data stream".to_owned())), + }; + + if res == Status::StreamEnd { + *eof = true; + } + Ok(vm.ctx.new_bytes(buf.to_vec())) + } + + #[pyproperty] + fn eof(&self) -> bool { + let state = self.state.lock(); + state.eof + } + + #[pyproperty] + fn unused_data(&self, vm: &VirtualMachine) -> PyBytesRef { + // Data found after the end of the compressed stream. + // If this attribute is accessed before the end of the stream + // has been reached, its value will be b''. + vm.ctx.new_bytes(b"".to_vec()) + // alternatively, be more honest: + // Err(vm.new_not_implemented_error( + // "unused_data isn't implemented yet".to_owned(), + // )) + // + // TODO + // let state = self.state.lock(); + // if state.eof { + // vm.ctx.new_bytes(state.input_buffer.to_vec()) + // else { + // vm.ctx.new_bytes(b"".to_vec()) + // } + } + + #[pyproperty] + fn needs_input(&self) -> bool { + // False if the decompress() method can provide more + // decompressed data before requiring new uncompressed input. + let state = self.state.lock(); + state.needs_input + } + + // TODO: mro()? + } + + struct CompressorState { + flushed: bool, + encoder: Option>>, + } + + #[pyattr] + #[pyclass(name = "BZ2Compressor")] + #[derive(PyPayload)] + struct BZ2Compressor { + state: PyMutex, + } + + impl fmt::Debug for BZ2Compressor { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "_bz2.BZ2Compressor") + } + } + + impl Constructor for BZ2Compressor { + type Args = (OptionalArg,); + + fn py_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult { + let (compresslevel,) = args; + // 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".to_owned()) + ) + } + }; + + Self { + state: PyMutex::new(CompressorState { + flushed: false, + encoder: Some(BzEncoder::new(Vec::new(), level)), + }), + } + .into_ref_with_type(vm, cls) + .map(Into::into) + } + } + + // TODO: return partial results from compress() instead of returning everything in flush() + #[pyclass(with(Constructor))] + impl BZ2Compressor { + #[pymethod] + fn compress(&self, data: ArgBytesLike, vm: &VirtualMachine) -> PyResult { + let mut state = self.state.lock(); + if state.flushed { + return Err(vm.new_value_error("Compressor has been flushed".to_owned())); + } + + // let CompressorState { flushed, encoder } = &mut *state; + let CompressorState { encoder, .. } = &mut *state; + + // TODO: handle Err + data.with_ref(|input_bytes| encoder.as_mut().unwrap().write_all(input_bytes).unwrap()); + Ok(vm.ctx.new_bytes(Vec::new())) + } + + #[pymethod] + fn flush(&self, vm: &VirtualMachine) -> PyResult { + let mut state = self.state.lock(); + if state.flushed { + return Err(vm.new_value_error("Repeated call to flush()".to_owned())); + } + + // 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())) + } + } +} diff --git a/stdlib/src/lib.rs b/stdlib/src/lib.rs index 98cffde6f..e610e0059 100644 --- a/stdlib/src/lib.rs +++ b/stdlib/src/lib.rs @@ -25,6 +25,8 @@ mod random; mod statistics; // TODO: maybe make this an extension module, if we ever get those // mod re; +#[cfg(not(any(target_arch = "wasm32")))] +mod bz2; #[cfg(not(target_arch = "wasm32"))] pub mod socket; #[cfg(unix)] @@ -115,6 +117,7 @@ pub fn get_module_inits() -> impl Iterator, StdlibInit "select" => select::make_module, "_socket" => socket::make_module, "faulthandler" => faulthandler::make_module, + "_bz2" => bz2::make_module, } #[cfg(feature = "ssl")] {