flatten compression modules

RustPython · youknowone · May 7, 2025 · May 7, 2025 · Apr 19, 2025 · Apr 19, 2025
commit d9c18c559380f1458b448cba43fd6b537dc576aa
diff --git a/stdlib/src/compression/bz2.rs → stdlib/src/bz2.rs b/stdlib/src/compression/bz2.rs → stdlib/src/bz2.rs
@@ -4,13 +4,13 @@ pub(crate) use _bz2::make_module;
 
 #[pymodule]
 mod _bz2 {
-    use crate::common::lock::PyMutex;
     use crate::compression::{
         DecompressArgs, DecompressError, DecompressState, DecompressStatus, Decompressor,
     };
     use crate::vm::{
         VirtualMachine,
         builtins::{PyBytesRef, PyTypeRef},
+        common::lock::PyMutex,
         function::{ArgBytesLike, OptionalArg},
         object::{PyPayload, PyResult},
         types::Constructor,

diff --git a/stdlib/src/compression/mod.rs → stdlib/src/compression.rs b/stdlib/src/compression/mod.rs → stdlib/src/compression.rs
@@ -1,20 +1,75 @@
-// spell-checker:ignore compressobj decompressobj zdict chunksize zlibmodule miniz chunker
+// cspell:ignore chunker
+
+//! internal shared module for compression libraries
 
 use crate::vm::function::{ArgBytesLike, ArgSize, OptionalArg};
 use crate::vm::{
     PyResult, VirtualMachine,
     builtins::{PyBaseExceptionRef, PyBytesRef},
     convert::ToPyException,
 };
-pub(crate) mod bz2;
-#[cfg(not(any(target_os = "android", target_arch = "wasm32")))]
-pub(crate) mod lzma;
-pub(crate) mod zlib;
 
 pub const USE_AFTER_FINISH_ERR: &str = "Error -2: inconsistent stream state";
 // TODO: don't hardcode
 const CHUNKSIZE: usize = u32::MAX as usize;
 
+#[derive(FromArgs)]
+pub struct DecompressArgs {
+    #[pyarg(positional)]
+    data: ArgBytesLike,
+    #[pyarg(any, optional)]
+    max_length: OptionalArg<ArgSize>,
+}
+
+impl DecompressArgs {
+    pub fn data(&self) -> crate::common::borrow::BorrowedValue<'_, [u8]> {
+        self.data.borrow_buf()
+    }
+    pub fn raw_max_length(&self) -> Option<isize> {
+        self.max_length.into_option().map(|ArgSize { value }| value)
+    }
+
+    // negative is None
+    pub fn max_length(&self) -> Option<usize> {
+        self.max_length
+            .into_option()
+            .and_then(|ArgSize { value }| usize::try_from(value).ok())
+    }
+}
+
+pub 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<u8>,
+        flush: Self::Flush,
+    ) -> Result<Self::Status, Self::Error>;
+    fn maybe_set_dict(&mut self, err: Self::Error) -> Result<(), Self::Error> {
+        Err(err)
+    }
+}
+
+pub trait DecompressStatus {
+    fn is_stream_end(&self) -> bool;
+}
+
+pub trait DecompressFlushKind: Copy {
+    const SYNC: Self;
+}
+
+impl DecompressFlushKind for () {
+    const SYNC: Self = ();
+}
+
+pub fn flush_sync<T: DecompressFlushKind>(_final_chunk: bool) -> T {
+    T::SYNC
+}
+
 #[derive(Clone)]
 pub struct Chunker<'a> {
     data1: &'a [u8],
@@ -57,6 +112,17 @@ impl<'a> Chunker<'a> {
     }
 }
 
+pub fn _decompress<D: Decompressor>(
+    data: &[u8],
+    d: &mut D,
+    bufsize: usize,
+    max_length: Option<usize>,
+    calc_flush: impl Fn(bool) -> D::Flush,
+) -> Result<(Vec<u8>, bool), D::Error> {
+    let mut data = Chunker::new(data);
+    _decompress_chunks(&mut data, d, bufsize, max_length, calc_flush)
+}
+
 pub fn _decompress_chunks<D: Decompressor>(
     data: &mut Chunker<'_>,
     d: &mut D,
@@ -207,39 +273,6 @@ impl<C: Compressor> CompressState<C> {
     }
 }
 
-pub 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<u8>,
-        flush: Self::Flush,
-    ) -> Result<Self::Status, Self::Error>;
-    fn maybe_set_dict(&mut self, err: Self::Error) -> Result<(), Self::Error> {
-        Err(err)
-    }
-}
-
-pub trait DecompressStatus {
-    fn is_stream_end(&self) -> bool;
-}
-
-pub trait DecompressFlushKind: Copy {
-    const SYNC: Self;
-}
-
-impl DecompressFlushKind for () {
-    const SYNC: Self = ();
-}
-
-pub fn flush_sync<T: DecompressFlushKind>(_final_chunk: bool) -> T {
-    T::SYNC
-}
-
 #[derive(Debug)]
 pub struct DecompressState<D> {
     decompress: D,
@@ -339,25 +372,3 @@ impl ToPyException for EofError {
         vm.new_eof_error("End of stream already reached".to_owned())
     }
 }
-
-#[derive(FromArgs)]
-pub struct DecompressArgs {
-    #[pyarg(positional)]
-    data: ArgBytesLike,
-    #[pyarg(any, optional)]
-    max_length: OptionalArg<ArgSize>,
-}
-
-impl DecompressArgs {
-    pub fn data(&self) -> crate::common::borrow::BorrowedValue<'_, [u8]> {
-        self.data.borrow_buf()
-    }
-    pub fn raw_max_length(&self) -> Option<isize> {
-        self.max_length.into_option().map(|ArgSize { value }| value)
-    }
-    pub fn max_length(&self) -> Option<usize> {
-        self.max_length
-            .into_option()
-            .and_then(|ArgSize { value }| usize::try_from(value).ok())
-    }
-}
diff --git a/stdlib/src/lib.rs b/stdlib/src/lib.rs
@@ -11,12 +11,17 @@ pub mod array;
 mod binascii;
 mod bisect;
 mod cmath;
-mod compression;
 mod contextvars;
 mod csv;
 mod dis;
 mod gc;
 
+mod bz2;
+mod compression; // internal module
+#[cfg(not(any(target_os = "android", target_arch = "wasm32")))]
+mod lzma;
+mod zlib;
+
 mod blake2;
 mod hashlib;
 mod md5;
@@ -111,7 +116,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" => compression::bz2::make_module,
+            "_bz2" => bz2::make_module,
             "cmath" => cmath::make_module,
             "_contextvars" => contextvars::make_module,
             "_csv" => csv::make_module,
@@ -132,7 +137,7 @@ pub fn get_module_inits() -> impl Iterator<Item = (Cow<'static, str>, StdlibInit
             "_statistics" => statistics::make_module,
             "_struct" => pystruct::make_module,
             "unicodedata" => unicodedata::make_module,
-            "zlib" => compression::zlib::make_module,
+            "zlib" => zlib::make_module,
             "_statistics" => statistics::make_module,
             "_suggestions" => suggestions::make_module,
             // crate::vm::sysmodule::sysconfigdata_name() => sysconfigdata::make_module,
@@ -152,7 +157,7 @@ pub fn get_module_inits() -> impl Iterator<Item = (Cow<'static, str>, StdlibInit
         }
         #[cfg(not(any(target_os = "android", target_arch = "wasm32")))]
         {
-            "_lzma" => compression::lzma::make_module,
+            "_lzma" => lzma::make_module,
         }
         #[cfg(all(feature = "sqlite", not(any(target_os = "android", target_arch = "wasm32"))))]
         {

diff --git a/stdlib/src/compression/lzma.rs → stdlib/src/lzma.rs b/stdlib/src/compression/lzma.rs → stdlib/src/lzma.rs
diff --git a/stdlib/src/compression/zlib.rs → stdlib/src/zlib.rs b/stdlib/src/compression/zlib.rs → stdlib/src/zlib.rs
@@ -5,9 +5,9 @@ pub(crate) use zlib::make_module;
 #[pymodule]
 mod zlib {
     use crate::compression::{
-        _decompress_chunks, Chunker, CompressFlushKind, CompressState, CompressStatusKind,
-        Compressor, DecompressArgs, DecompressError, DecompressFlushKind, DecompressState,
-        DecompressStatus, Decompressor, USE_AFTER_FINISH_ERR, flush_sync,
+        _decompress, CompressFlushKind, CompressState, CompressStatusKind, Compressor,
+        DecompressArgs, DecompressError, DecompressFlushKind, DecompressState, DecompressStatus,
+        Decompressor, USE_AFTER_FINISH_ERR, flush_sync,
     };
     use crate::vm::{
         PyObject, PyPayload, PyResult, VirtualMachine,
@@ -146,17 +146,6 @@ mod zlib {
         }
     }
 
-    fn _decompress<D: Decompressor>(
-        data: &[u8],
-        d: &mut D,
-        bufsize: usize,
-        max_length: Option<usize>,
-        calc_flush: impl Fn(bool) -> D::Flush,
-    ) -> Result<(Vec<u8>, bool), D::Error> {
-        let mut data = Chunker::new(data);
-        _decompress_chunks(&mut data, d, bufsize, max_length, calc_flush)
-    }
-
     #[derive(FromArgs)]
     struct PyFuncDecompressArgs {
         #[pyarg(positional)]