RustPython
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -19,7 +19,6 @@ flame-it = ["rustpython-vm/flame-it", "flame", "flamescope"]
 freeze-stdlib = ["stdlib", "rustpython-vm/freeze-stdlib", "rustpython-pylib?/freeze-stdlib"]
 jit = ["rustpython-vm/jit"]
 threading = ["rustpython-vm/threading", "rustpython-stdlib/threading"]
-bz2 = ["stdlib", "rustpython-stdlib/bz2"]
 sqlite = ["rustpython-stdlib/sqlite"]
 ssl = ["rustpython-stdlib/ssl"]
 ssl-vendor = ["ssl", "rustpython-stdlib/ssl-vendor"]

diff --git a/Lib/test/test_bz2.py b/Lib/test/test_bz2.py
@@ -676,6 +676,8 @@ def testCompress4G(self, size):
         finally:
             data = None
 
+    # TODO: RUSTPYTHON
+    @unittest.expectedFailure
     def testPickle(self):
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             with self.assertRaises(TypeError):
@@ -734,6 +736,8 @@ def testDecompress4G(self, size):
             compressed = None
             decompressed = None
 
+    # TODO: RUSTPYTHON
+    @unittest.expectedFailure
     def testPickle(self):
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             with self.assertRaises(TypeError):
@@ -1001,6 +1005,8 @@ def test_encoding_error_handler(self):
                 as f:
             self.assertEqual(f.read(), "foobar")
 
+    # TODO: RUSTPYTHON
+    @unittest.expectedFailure
     def test_newline(self):
         # Test with explicit newline (universal newline mode disabled).
         text = self.TEXT.decode("ascii")

diff --git a/stdlib/Cargo.toml b/stdlib/Cargo.toml
@@ -14,7 +14,6 @@ license.workspace = true
 default = ["compiler"]
 compiler = ["rustpython-vm/compiler"]
 threading = ["rustpython-common/threading", "rustpython-vm/threading"]
-bz2 = ["bzip2"]
 sqlite = ["dep:libsqlite3-sys"]
 ssl = ["openssl", "openssl-sys", "foreign-types-shared", "openssl-probe"]
 ssl-vendor = ["ssl", "openssl/vendored"]
@@ -80,7 +79,7 @@ adler32 = "1.2.0"
 crc32fast = "1.3.2"
 flate2 = { version = "1.1", default-features = false, features = ["zlib-rs"] }
 libz-sys = { package = "libz-rs-sys", version = "0.5" }
-bzip2 = { version = "0.4", optional = true }
+bzip2 = { version = "0.5", features = ["libbz2-rs-sys"] }
 
 # tkinter
 tk-sys = { git = "https://github.com/arihant2math/tkinter.git", tag = "v0.1.0", optional = true }

diff --git a/stdlib/src/bz2.rs b/stdlib/src/bz2.rs
@@ -12,28 +12,48 @@ mod _bz2 {
         object::{PyPayload, PyResult},
         types::Constructor,
     };
+    use crate::zlib::{
+        DecompressArgs, DecompressError, DecompressState, DecompressStatus, Decompressor,
+    };
     use bzip2::{Decompress, Status, write::BzEncoder};
+    use rustpython_vm::convert::ToPyException;
     use std::{fmt, io::Write};
 
-    // const BUFSIZ: i32 = 8192;
-
-    struct DecompressorState {
-        decoder: Decompress,
-        eof: bool,
-        needs_input: bool,
-        // input_buffer: Vec<u8>,
-        // output_buffer: Vec<u8>,
-    }
+    const BUFSIZ: usize = 8192;
 
     #[pyattr]
     #[pyclass(name = "BZ2Decompressor")]
     #[derive(PyPayload)]
     struct BZ2Decompressor {
-        state: PyMutex<DecompressorState>,
+        state: PyMutex<DecompressState<Decompress>>,
+    }
+
+    impl Decompressor for Decompress {
+        type Flush = ();
+        type Status = Status;
+        type Error = bzip2::Error;
+
+        fn total_in(&self) -> u64 {
+            self.total_in()
+        }
+        fn decompress_vec(
+            &mut self,
+            input: &[u8],
+            output: &mut Vec<u8>,
+            (): Self::Flush,
+        ) -> Result<Self::Status, Self::Error> {
+            self.decompress_vec(input, output)
+        }
+    }
+
+    impl DecompressStatus for Status {
+        fn is_stream_end(&self) -> bool {
+            *self == Status::StreamEnd
+        }
     }
 
     impl fmt::Debug for BZ2Decompressor {
-        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
             write!(f, "_bz2.BZ2Decompressor")
         }
     }
@@ -43,13 +63,7 @@ mod _bz2 {
 
         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(),
-                }),
+                state: PyMutex::new(DecompressState::new(Decompress::new(false), vm)),
             }
             .into_ref_with_type(vm, cls)
             .map(Into::into)
@@ -59,107 +73,34 @@ mod _bz2 {
     #[pyclass(with(Constructor))]
     impl BZ2Decompressor {
         #[pymethod]
-        fn decompress(
-            &self,
-            data: ArgBytesLike,
-            // TODO: PyIntRef
-            max_length: OptionalArg<i32>,
-            vm: &VirtualMachine,
-        ) -> PyResult<PyBytesRef> {
-            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
-            // };
+        fn decompress(&self, args: DecompressArgs, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
+            let max_length = args.max_length();
+            let data = &*args.data();
 
             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()))
+            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),
+                })
         }
 
         #[pygetset]
         fn eof(&self) -> bool {
-            let state = self.state.lock();
-            state.eof
+            self.state.lock().eof()
         }
 
         #[pygetset]
-        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())
-            // }
+        fn unused_data(&self) -> PyBytesRef {
+            self.state.lock().unused_data()
         }
 
         #[pygetset]
         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
+            self.state.lock().needs_input()
         }
 
         // TODO: mro()?
@@ -178,7 +119,7 @@ mod _bz2 {
     }
 
     impl fmt::Debug for BZ2Compressor {
-        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
             write!(f, "_bz2.BZ2Compressor")
         }
     }

diff --git a/stdlib/src/lib.rs b/stdlib/src/lib.rs
@@ -36,7 +36,6 @@ mod statistics;
 mod suggestions;
 // TODO: maybe make this an extension module, if we ever get those
 // mod re;
-#[cfg(feature = "bz2")]
 mod bz2;
 #[cfg(not(target_arch = "wasm32"))]
 pub mod socket;
@@ -112,6 +111,7 @@ pub fn get_module_inits() -> impl Iterator<Item = (Cow<'static, str>, StdlibInit
             "array" => array::make_module,
             "binascii" => binascii::make_module,
             "_bisect" => bisect::make_module,
+            "_bz2" => bz2::make_module,
             "cmath" => cmath::make_module,
             "_contextvars" => contextvars::make_module,
             "_csv" => csv::make_module,
@@ -158,10 +158,6 @@ pub fn get_module_inits() -> impl Iterator<Item = (Cow<'static, str>, StdlibInit
         {
             "_ssl" => ssl::make_module,
         }
-        #[cfg(feature = "bz2")]
-        {
-            "_bz2" => bz2::make_module,
-        }
         #[cfg(windows)]
         {
             "_overlapped" => overlapped::make_module,