Skip to content

Commit 663ce18

Browse files
committed
flatten compression modules
1 parent acae154 commit 663ce18

File tree

5 files changed

+83
-79
lines changed

5 files changed

+83
-79
lines changed

stdlib/src/compression/bz2.rs renamed to stdlib/src/bz2.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,13 @@ pub(crate) use _bz2::make_module;
44

55
#[pymodule]
66
mod _bz2 {
7-
use crate::common::lock::PyMutex;
87
use crate::compression::{
98
DecompressArgs, DecompressError, DecompressState, DecompressStatus, Decompressor,
109
};
1110
use crate::vm::{
1211
VirtualMachine,
1312
builtins::{PyBytesRef, PyTypeRef},
13+
common::lock::PyMutex,
1414
function::{ArgBytesLike, OptionalArg},
1515
object::{PyPayload, PyResult},
1616
types::Constructor,

stdlib/src/compression/mod.rs renamed to stdlib/src/compression.rs

+71-60
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,75 @@
1-
// spell-checker:ignore compressobj decompressobj zdict chunksize zlibmodule miniz chunker
1+
// cspell:ignore chunker
2+
3+
//! internal shared module for compression libraries
24
35
use crate::vm::function::{ArgBytesLike, ArgSize, OptionalArg};
46
use crate::vm::{
57
PyResult, VirtualMachine,
68
builtins::{PyBaseExceptionRef, PyBytesRef},
79
convert::ToPyException,
810
};
9-
pub(crate) mod bz2;
10-
#[cfg(not(any(target_os = "android", target_arch = "wasm32")))]
11-
pub(crate) mod lzma;
12-
pub(crate) mod zlib;
1311

1412
pub const USE_AFTER_FINISH_ERR: &str = "Error -2: inconsistent stream state";
1513
// TODO: don't hardcode
1614
const CHUNKSIZE: usize = u32::MAX as usize;
1715

16+
#[derive(FromArgs)]
17+
pub struct DecompressArgs {
18+
#[pyarg(positional)]
19+
data: ArgBytesLike,
20+
#[pyarg(any, optional)]
21+
max_length: OptionalArg<ArgSize>,
22+
}
23+
24+
impl DecompressArgs {
25+
pub fn data(&self) -> crate::common::borrow::BorrowedValue<'_, [u8]> {
26+
self.data.borrow_buf()
27+
}
28+
pub fn raw_max_length(&self) -> Option<isize> {
29+
self.max_length.into_option().map(|ArgSize { value }| value)
30+
}
31+
32+
// negative is None
33+
pub fn max_length(&self) -> Option<usize> {
34+
self.max_length
35+
.into_option()
36+
.and_then(|ArgSize { value }| usize::try_from(value).ok())
37+
}
38+
}
39+
40+
pub trait Decompressor {
41+
type Flush: DecompressFlushKind;
42+
type Status: DecompressStatus;
43+
type Error;
44+
45+
fn total_in(&self) -> u64;
46+
fn decompress_vec(
47+
&mut self,
48+
input: &[u8],
49+
output: &mut Vec<u8>,
50+
flush: Self::Flush,
51+
) -> Result<Self::Status, Self::Error>;
52+
fn maybe_set_dict(&mut self, err: Self::Error) -> Result<(), Self::Error> {
53+
Err(err)
54+
}
55+
}
56+
57+
pub trait DecompressStatus {
58+
fn is_stream_end(&self) -> bool;
59+
}
60+
61+
pub trait DecompressFlushKind: Copy {
62+
const SYNC: Self;
63+
}
64+
65+
impl DecompressFlushKind for () {
66+
const SYNC: Self = ();
67+
}
68+
69+
pub fn flush_sync<T: DecompressFlushKind>(_final_chunk: bool) -> T {
70+
T::SYNC
71+
}
72+
1873
#[derive(Clone)]
1974
pub struct Chunker<'a> {
2075
data1: &'a [u8],
@@ -57,6 +112,17 @@ impl<'a> Chunker<'a> {
57112
}
58113
}
59114

115+
pub fn _decompress<D: Decompressor>(
116+
data: &[u8],
117+
d: &mut D,
118+
bufsize: usize,
119+
max_length: Option<usize>,
120+
calc_flush: impl Fn(bool) -> D::Flush,
121+
) -> Result<(Vec<u8>, bool), D::Error> {
122+
let mut data = Chunker::new(data);
123+
_decompress_chunks(&mut data, d, bufsize, max_length, calc_flush)
124+
}
125+
60126
pub fn _decompress_chunks<D: Decompressor>(
61127
data: &mut Chunker<'_>,
62128
d: &mut D,
@@ -207,39 +273,6 @@ impl<C: Compressor> CompressState<C> {
207273
}
208274
}
209275

210-
pub trait Decompressor {
211-
type Flush: DecompressFlushKind;
212-
type Status: DecompressStatus;
213-
type Error;
214-
215-
fn total_in(&self) -> u64;
216-
fn decompress_vec(
217-
&mut self,
218-
input: &[u8],
219-
output: &mut Vec<u8>,
220-
flush: Self::Flush,
221-
) -> Result<Self::Status, Self::Error>;
222-
fn maybe_set_dict(&mut self, err: Self::Error) -> Result<(), Self::Error> {
223-
Err(err)
224-
}
225-
}
226-
227-
pub trait DecompressStatus {
228-
fn is_stream_end(&self) -> bool;
229-
}
230-
231-
pub trait DecompressFlushKind: Copy {
232-
const SYNC: Self;
233-
}
234-
235-
impl DecompressFlushKind for () {
236-
const SYNC: Self = ();
237-
}
238-
239-
pub fn flush_sync<T: DecompressFlushKind>(_final_chunk: bool) -> T {
240-
T::SYNC
241-
}
242-
243276
#[derive(Debug)]
244277
pub struct DecompressState<D> {
245278
decompress: D,
@@ -339,25 +372,3 @@ impl ToPyException for EofError {
339372
vm.new_eof_error("End of stream already reached".to_owned())
340373
}
341374
}
342-
343-
#[derive(FromArgs)]
344-
pub struct DecompressArgs {
345-
#[pyarg(positional)]
346-
data: ArgBytesLike,
347-
#[pyarg(any, optional)]
348-
max_length: OptionalArg<ArgSize>,
349-
}
350-
351-
impl DecompressArgs {
352-
pub fn data(&self) -> crate::common::borrow::BorrowedValue<'_, [u8]> {
353-
self.data.borrow_buf()
354-
}
355-
pub fn raw_max_length(&self) -> Option<isize> {
356-
self.max_length.into_option().map(|ArgSize { value }| value)
357-
}
358-
pub fn max_length(&self) -> Option<usize> {
359-
self.max_length
360-
.into_option()
361-
.and_then(|ArgSize { value }| usize::try_from(value).ok())
362-
}
363-
}

stdlib/src/lib.rs

+8-4
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,16 @@ pub mod array;
1111
mod binascii;
1212
mod bisect;
1313
mod cmath;
14-
mod compression;
1514
mod contextvars;
1615
mod csv;
1716
mod dis;
1817
mod gc;
1918

19+
mod compression; // internal module
20+
mod lzma;
21+
mod zlib;
22+
mod bz2;
23+
2024
mod blake2;
2125
mod hashlib;
2226
mod md5;
@@ -111,7 +115,7 @@ pub fn get_module_inits() -> impl Iterator<Item = (Cow<'static, str>, StdlibInit
111115
"array" => array::make_module,
112116
"binascii" => binascii::make_module,
113117
"_bisect" => bisect::make_module,
114-
"_bz2" => compression::bz2::make_module,
118+
"_bz2" => bz2::make_module,
115119
"cmath" => cmath::make_module,
116120
"_contextvars" => contextvars::make_module,
117121
"_csv" => csv::make_module,
@@ -132,7 +136,7 @@ pub fn get_module_inits() -> impl Iterator<Item = (Cow<'static, str>, StdlibInit
132136
"_statistics" => statistics::make_module,
133137
"_struct" => pystruct::make_module,
134138
"unicodedata" => unicodedata::make_module,
135-
"zlib" => compression::zlib::make_module,
139+
"zlib" => zlib::make_module,
136140
"_statistics" => statistics::make_module,
137141
"_suggestions" => suggestions::make_module,
138142
// crate::vm::sysmodule::sysconfigdata_name() => sysconfigdata::make_module,
@@ -152,7 +156,7 @@ pub fn get_module_inits() -> impl Iterator<Item = (Cow<'static, str>, StdlibInit
152156
}
153157
#[cfg(not(any(target_os = "android", target_arch = "wasm32")))]
154158
{
155-
"_lzma" => compression::lzma::make_module,
159+
"_lzma" => lzma::make_module,
156160
}
157161
#[cfg(all(feature = "sqlite", not(any(target_os = "android", target_arch = "wasm32"))))]
158162
{
File renamed without changes.

stdlib/src/compression/zlib.rs renamed to stdlib/src/zlib.rs

+3-14
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,9 @@ pub(crate) use zlib::make_module;
55
#[pymodule]
66
mod zlib {
77
use crate::compression::{
8-
_decompress_chunks, Chunker, CompressFlushKind, CompressState, CompressStatusKind,
9-
Compressor, DecompressArgs, DecompressError, DecompressFlushKind, DecompressState,
10-
DecompressStatus, Decompressor, USE_AFTER_FINISH_ERR, flush_sync,
8+
_decompress, CompressFlushKind, CompressState, CompressStatusKind, Compressor,
9+
DecompressArgs, DecompressError, DecompressFlushKind, DecompressState, DecompressStatus,
10+
Decompressor, USE_AFTER_FINISH_ERR, flush_sync,
1111
};
1212
use crate::vm::{
1313
PyObject, PyPayload, PyResult, VirtualMachine,
@@ -146,17 +146,6 @@ mod zlib {
146146
}
147147
}
148148

149-
fn _decompress<D: Decompressor>(
150-
data: &[u8],
151-
d: &mut D,
152-
bufsize: usize,
153-
max_length: Option<usize>,
154-
calc_flush: impl Fn(bool) -> D::Flush,
155-
) -> Result<(Vec<u8>, bool), D::Error> {
156-
let mut data = Chunker::new(data);
157-
_decompress_chunks(&mut data, d, bufsize, max_length, calc_flush)
158-
}
159-
160149
#[derive(FromArgs)]
161150
struct PyFuncDecompressArgs {
162151
#[pyarg(positional)]

0 commit comments

Comments
 (0)