@@ -6,7 +6,8 @@ pub(crate) use _bz2::make_module;
6
6
mod _bz2 {
7
7
use crate :: common:: lock:: PyMutex ;
8
8
use crate :: vm:: {
9
- FromArgs , VirtualMachine ,
9
+ FromArgs ,
10
+ VirtualMachine ,
10
11
builtins:: { PyBytesRef , PyTypeRef } ,
11
12
function:: { ArgBytesLike , OptionalArg } ,
12
13
object:: { PyPayload , PyResult } ,
@@ -25,7 +26,7 @@ mod _bz2 {
25
26
eof : bool ,
26
27
// Unused data found after the end of stream.
27
28
unused_data : Option < Vec < u8 > > ,
28
- needs_input : bool ,
29
+ needs_input : bool
29
30
}
30
31
31
32
#[ pyattr]
@@ -50,7 +51,7 @@ mod _bz2 {
50
51
eof : false ,
51
52
input_buffer : Vec :: new ( ) ,
52
53
unused_data : None ,
53
- needs_input : true ,
54
+ needs_input : true
54
55
} ) ,
55
56
}
56
57
. into_ref_with_type ( vm, cls)
@@ -62,27 +63,25 @@ mod _bz2 {
62
63
struct DecompressArgs {
63
64
#[ pyarg( positional) ]
64
65
data : ArgBytesLike ,
65
- #[ pyarg( any, default = "-1" ) ]
66
+ #[ pyarg( any, default = - 1 ) ]
66
67
max_length : i64 ,
67
68
}
68
69
69
70
#[ pyclass( with( Constructor ) ) ]
70
71
impl BZ2Decompressor {
71
72
#[ pymethod]
72
- fn decompress ( & self , args : DecompressArgs , vm : & VirtualMachine ) -> PyResult < PyBytesRef > {
73
+ fn decompress (
74
+ & self ,
75
+ args : DecompressArgs ,
76
+ vm : & VirtualMachine ,
77
+ ) -> PyResult < PyBytesRef > {
73
78
let DecompressArgs { data, max_length } = args;
74
79
let DecompressorState {
75
80
eof,
76
81
input_buffer,
77
82
unused_data,
78
83
needs_input,
79
84
} = & mut * self . state . lock ( ) ;
80
- if * eof {
81
- return Err ( vm. new_exception_msg (
82
- vm. ctx . exceptions . eof_error . to_owned ( ) ,
83
- "End of stream already reached" . to_owned ( ) ,
84
- ) ) ;
85
- }
86
85
let data_vec = data. borrow_buf ( ) . to_vec ( ) ;
87
86
input_buffer. extend ( data_vec) ;
88
87
@@ -95,35 +94,33 @@ mod _bz2 {
95
94
// If max_length is nonnegative, read at most that many bytes.
96
95
if max_length >= 0 {
97
96
let mut limited = decoder. by_ref ( ) . take ( max_length as u64 ) ;
98
- limited
99
- . read_to_end ( & mut output )
100
- . map_err ( |e| vm . new_os_error ( format ! ( "Decompression error: {}" , e ) ) ) ?;
97
+ limited. read_to_end ( & mut output ) . map_err ( |e| {
98
+ vm . new_os_error ( format ! ( "Decompression error: {}" , e ) )
99
+ } ) ?;
101
100
} else {
102
- decoder
103
- . read_to_end ( & mut output )
104
- . map_err ( |e| vm . new_os_error ( format ! ( "Decompression error: {}" , e ) ) ) ?;
101
+ decoder. read_to_end ( & mut output ) . map_err ( |e| {
102
+ vm . new_os_error ( format ! ( "Decompression error: {}" , e ) )
103
+ } ) ?;
105
104
}
106
105
107
106
// Determine how many bytes were consumed from the input.
108
107
let consumed = cursor. position ( ) as usize ;
109
108
// Remove the consumed bytes.
110
109
input_buffer. drain ( 0 ..consumed) ;
111
- unused_data. replace ( input_buffer. clone ( ) ) ;
112
- // skrink the vector to save memory
113
- input_buffer. shrink_to_fit ( ) ;
114
- if let Some ( v) = unused_data. as_mut ( ) {
115
- v. shrink_to_fit ( ) ;
116
- }
117
110
118
111
if * eof {
119
112
* needs_input = false ;
120
113
} else {
121
114
* needs_input = input_buffer. is_empty ( ) ;
122
115
}
116
+ let data_vec = data. borrow_buf ( ) . to_vec ( ) ;
117
+ input_buffer. extend ( data_vec) ;
123
118
124
119
// If the decoder reached end-of-stream (i.e. no more input remains), mark eof.
125
120
if input_buffer. is_empty ( ) {
126
121
* eof = true ;
122
+ * unused_data = Some ( input_buffer. clone ( ) ) ;
123
+ input_buffer. clear ( ) ;
127
124
}
128
125
129
126
Ok ( vm. ctx . new_bytes ( output) )
@@ -138,9 +135,10 @@ mod _bz2 {
138
135
#[ pygetset]
139
136
fn unused_data ( & self , vm : & VirtualMachine ) -> PyBytesRef {
140
137
let state = self . state . lock ( ) ;
141
- match & state. unused_data {
142
- Some ( data) => vm. ctx . new_bytes ( data. clone ( ) ) ,
143
- None => vm. ctx . new_bytes ( Vec :: new ( ) ) ,
138
+ if state. eof {
139
+ vm. ctx . new_bytes ( state. input_buffer . to_vec ( ) )
140
+ } else {
141
+ vm. ctx . new_bytes ( b"" . to_vec ( ) )
144
142
}
145
143
}
146
144
0 commit comments