@@ -12,28 +12,48 @@ mod _bz2 {
1212 object:: { PyPayload , PyResult } ,
1313 types:: Constructor ,
1414 } ;
15+ use crate :: zlib:: {
16+ DecompressArgs , DecompressError , DecompressState , DecompressStatus , Decompressor ,
17+ } ;
1518 use bzip2:: { Decompress , Status , write:: BzEncoder } ;
19+ use rustpython_vm:: convert:: ToPyException ;
1620 use std:: { fmt, io:: Write } ;
1721
18- // const BUFSIZ: i32 = 8192;
19-
20- struct DecompressorState {
21- decoder : Decompress ,
22- eof : bool ,
23- needs_input : bool ,
24- // input_buffer: Vec<u8>,
25- // output_buffer: Vec<u8>,
26- }
22+ const BUFSIZ : usize = 8192 ;
2723
2824 #[ pyattr]
2925 #[ pyclass( name = "BZ2Decompressor" ) ]
3026 #[ derive( PyPayload ) ]
3127 struct BZ2Decompressor {
32- state : PyMutex < DecompressorState > ,
28+ state : PyMutex < DecompressState < Decompress > > ,
29+ }
30+
31+ impl Decompressor for Decompress {
32+ type Flush = ( ) ;
33+ type Status = Status ;
34+ type Error = bzip2:: Error ;
35+
36+ fn total_in ( & self ) -> u64 {
37+ self . total_in ( )
38+ }
39+ fn decompress_vec (
40+ & mut self ,
41+ input : & [ u8 ] ,
42+ output : & mut Vec < u8 > ,
43+ ( ) : Self :: Flush ,
44+ ) -> Result < Self :: Status , Self :: Error > {
45+ self . decompress_vec ( input, output)
46+ }
47+ }
48+
49+ impl DecompressStatus for Status {
50+ fn is_stream_end ( & self ) -> bool {
51+ * self == Status :: StreamEnd
52+ }
3353 }
3454
3555 impl fmt:: Debug for BZ2Decompressor {
36- fn fmt ( & self , f : & mut fmt:: Formatter ) -> fmt:: Result {
56+ fn fmt ( & self , f : & mut fmt:: Formatter < ' _ > ) -> fmt:: Result {
3757 write ! ( f, "_bz2.BZ2Decompressor" )
3858 }
3959 }
@@ -43,13 +63,7 @@ mod _bz2 {
4363
4464 fn py_new ( cls : PyTypeRef , _: Self :: Args , vm : & VirtualMachine ) -> PyResult {
4565 Self {
46- state : PyMutex :: new ( DecompressorState {
47- decoder : Decompress :: new ( false ) ,
48- eof : false ,
49- needs_input : true ,
50- // input_buffer: Vec::new(),
51- // output_buffer: Vec::new(),
52- } ) ,
66+ state : PyMutex :: new ( DecompressState :: new ( Decompress :: new ( false ) , vm) ) ,
5367 }
5468 . into_ref_with_type ( vm, cls)
5569 . map ( Into :: into)
@@ -59,107 +73,34 @@ mod _bz2 {
5973 #[ pyclass( with( Constructor ) ) ]
6074 impl BZ2Decompressor {
6175 #[ pymethod]
62- fn decompress (
63- & self ,
64- data : ArgBytesLike ,
65- // TODO: PyIntRef
66- max_length : OptionalArg < i32 > ,
67- vm : & VirtualMachine ,
68- ) -> PyResult < PyBytesRef > {
69- let max_length = max_length. unwrap_or ( -1 ) ;
70- if max_length >= 0 {
71- return Err ( vm. new_not_implemented_error (
72- "the max_value argument is not implemented yet" . to_owned ( ) ,
73- ) ) ;
74- }
75- // let max_length = if max_length < 0 || max_length >= BUFSIZ {
76- // BUFSIZ
77- // } else {
78- // max_length
79- // };
76+ fn decompress ( & self , args : DecompressArgs , vm : & VirtualMachine ) -> PyResult < Vec < u8 > > {
77+ let max_length = args. max_length ( ) ;
78+ let data = & * args. data ( ) ;
8079
8180 let mut state = self . state . lock ( ) ;
82- let DecompressorState {
83- decoder,
84- eof,
85- ..
86- // needs_input,
87- // input_buffer,
88- // output_buffer,
89- } = & mut * state;
90-
91- if * eof {
92- return Err ( vm. new_exception_msg (
93- vm. ctx . exceptions . eof_error . to_owned ( ) ,
94- "End of stream already reached" . to_owned ( ) ,
95- ) ) ;
96- }
97-
98- // data.with_ref(|data| input_buffer.extend(data));
99-
100- // If max_length is negative:
101- // read the input X bytes at a time, compress it and append it to output.
102- // Once you're out of input, setting needs_input to true and return the
103- // output as bytes.
104- //
105- // TODO:
106- // If max_length is non-negative:
107- // Read the input X bytes at a time, compress it and append it to
108- // the output. If output reaches `max_length` in size, return
109- // it (up to max_length), and store the rest of the output
110- // for later.
111-
112- // TODO: arbitrary choice, not the right way to do it.
113- let mut buf = Vec :: with_capacity ( data. len ( ) * 32 ) ;
114-
115- let before = decoder. total_in ( ) ;
116- let res = data. with_ref ( |data| decoder. decompress_vec ( data, & mut buf) ) ;
117- let _written = ( decoder. total_in ( ) - before) as usize ;
118-
119- let res = match res {
120- Ok ( x) => x,
121- // TODO: error message
122- _ => return Err ( vm. new_os_error ( "Invalid data stream" . to_owned ( ) ) ) ,
123- } ;
124-
125- if res == Status :: StreamEnd {
126- * eof = true ;
127- }
128- Ok ( vm. ctx . new_bytes ( buf. to_vec ( ) ) )
81+ state
82+ . decompress ( data, max_length, BUFSIZ , vm)
83+ . map_err ( |e| match e {
84+ DecompressError :: Decompress ( err) => vm. new_os_error ( err. to_string ( ) ) ,
85+ DecompressError :: Eof ( err) => err. to_pyexception ( vm) ,
86+ } )
12987 }
13088
13189 #[ pygetset]
13290 fn eof ( & self ) -> bool {
133- let state = self . state . lock ( ) ;
134- state. eof
91+ self . state . lock ( ) . eof ( )
13592 }
13693
13794 #[ pygetset]
138- fn unused_data ( & self , vm : & VirtualMachine ) -> PyBytesRef {
139- // Data found after the end of the compressed stream.
140- // If this attribute is accessed before the end of the stream
141- // has been reached, its value will be b''.
142- vm. ctx . new_bytes ( b"" . to_vec ( ) )
143- // alternatively, be more honest:
144- // Err(vm.new_not_implemented_error(
145- // "unused_data isn't implemented yet".to_owned(),
146- // ))
147- //
148- // TODO
149- // let state = self.state.lock();
150- // if state.eof {
151- // vm.ctx.new_bytes(state.input_buffer.to_vec())
152- // else {
153- // vm.ctx.new_bytes(b"".to_vec())
154- // }
95+ fn unused_data ( & self ) -> PyBytesRef {
96+ self . state . lock ( ) . unused_data ( )
15597 }
15698
15799 #[ pygetset]
158100 fn needs_input ( & self ) -> bool {
159101 // False if the decompress() method can provide more
160102 // decompressed data before requiring new uncompressed input.
161- let state = self . state . lock ( ) ;
162- state. needs_input
103+ self . state . lock ( ) . needs_input ( )
163104 }
164105
165106 // TODO: mro()?
@@ -178,7 +119,7 @@ mod _bz2 {
178119 }
179120
180121 impl fmt:: Debug for BZ2Compressor {
181- fn fmt ( & self , f : & mut fmt:: Formatter ) -> fmt:: Result {
122+ fn fmt ( & self , f : & mut fmt:: Formatter < ' _ > ) -> fmt:: Result {
182123 write ! ( f, "_bz2.BZ2Compressor" )
183124 }
184125 }
0 commit comments