From 691ee06dee6913f83b02415c42d75bccd98da439 Mon Sep 17 00:00:00 2001
From: Kwak Byoung Min <67214970+doma17@users.noreply.github.com>
Date: Tue, 21 Jul 2026 00:30:10 +0900
Subject: [PATCH] bz2: return streaming output from BZ2Compressor.compress()
 (#8330)

* bz2: return streaming output from BZ2Compressor.compress()

  Return output produced during compression instead of waiting for flush().

  Constraint: Preserve encoder state without flushing the BZ2 stream
  Confidence: high
  Scope-risk: narrow
  Directive: Do not flush merely to retrieve partial output
  Tested: test_bz2; stdlib_bz2.py; prek run --all-files
  Assisted-by: Codex:gpt-5.6-sol

* bz2: remove redundant streaming-output snippet
---
 Lib/test/test_bz2.py     |  1 -
 crates/stdlib/src/bz2.rs | 12 ++++++------
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/Lib/test/test_bz2.py b/Lib/test/test_bz2.py
index a7e152fb7..bcd4e033b 100644
--- a/Lib/test/test_bz2.py
+++ b/Lib/test/test_bz2.py
@@ -936,7 +936,6 @@ class BZ2DecompressorTest(BaseTest):
             with self.assertRaises(TypeError):
                 pickle.dumps(BZ2Decompressor(), proto)
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: 0 != 100
     def testDecompressorChunksMaxsize(self):
         bzd = BZ2Decompressor()
         max_length = 100
diff --git a/crates/stdlib/src/bz2.rs b/crates/stdlib/src/bz2.rs
index 575f33c4b..fcde54057 100644
--- a/crates/stdlib/src/bz2.rs
+++ b/crates/stdlib/src/bz2.rs
@@ -17,6 +17,7 @@ mod _bz2 {
     };
     use alloc::fmt;
     use bzip2::{Decompress, Status, write::BzEncoder};
+    use core::mem;
     use rustpython_vm::convert::ToPyException;
     use std::io::Write;
 
@@ -155,7 +156,6 @@ mod _bz2 {
         }
     }
 
-    // TODO: return partial results from compress() instead of returning everything in flush()
     #[pyclass(with(Constructor))]
     impl BZ2Compressor {
         #[pymethod]
@@ -165,12 +165,12 @@ mod _bz2 {
                 return Err(vm.new_value_error("Compressor has been flushed"));
             }
 
-            // 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()))
+            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())))
         }
 
         #[pymethod]