@@ -9,24 +9,22 @@ mod decl {
9
9
} ,
10
10
bytecode,
11
11
convert:: ToPyObject ,
12
- function:: ArgBytesLike ,
12
+ function:: { ArgBytesLike , OptionalArg } ,
13
13
object:: AsObject ,
14
14
protocol:: PyBuffer ,
15
15
PyObjectRef , PyResult , TryFromObject , VirtualMachine ,
16
16
} ;
17
- /// TODO
18
- /// PyBytes: Currently getting recursion error with match_class!
19
17
use num_bigint:: { BigInt , Sign } ;
20
18
use num_traits:: Zero ;
21
19
22
20
#[ repr( u8 ) ]
23
21
enum Type {
24
22
// Null = b'0',
25
- // None = b'N',
23
+ None = b'N' ,
26
24
False = b'F' ,
27
25
True = b'T' ,
28
26
// StopIter = b'S',
29
- // Ellipsis = b'.',
27
+ Ellipsis = b'.' ,
30
28
Int = b'i' ,
31
29
Float = b'g' ,
32
30
// Complex = b'y',
@@ -38,11 +36,11 @@ mod decl {
38
36
List = b'[' ,
39
37
Dict = b'{' ,
40
38
Code = b'c' ,
41
- Str = b'u' , // = TYPE_UNICODE
39
+ Unicode = b'u' ,
42
40
// Unknown = b'?',
43
41
Set = b'<' ,
44
42
FrozenSet = b'>' ,
45
- // Ascii = b'a',
43
+ Ascii = b'a' ,
46
44
// AsciiInterned = b'A',
47
45
// SmallTuple = b')',
48
46
// ShortAscii = b'z',
@@ -56,11 +54,11 @@ mod decl {
56
54
use Type :: * ;
57
55
Ok ( match value {
58
56
// b'0' => Null,
59
- // b'N' => None,
57
+ b'N' => None ,
60
58
b'F' => False ,
61
59
b'T' => True ,
62
60
// b'S' => StopIter,
63
- // b'.' => Ellipsis,
61
+ b'.' => Ellipsis ,
64
62
b'i' => Int ,
65
63
b'g' => Float ,
66
64
// b'y' => Complex,
@@ -72,11 +70,11 @@ mod decl {
72
70
b'[' => List ,
73
71
b'{' => Dict ,
74
72
b'c' => Code ,
75
- b'u' => Str ,
73
+ b'u' => Unicode ,
76
74
// b'?' => Unknown,
77
75
b'<' => Set ,
78
76
b'>' => FrozenSet ,
79
- // b'a' => Ascii,
77
+ b'a' => Ascii ,
80
78
// b'A' => AsciiInterned,
81
79
// b')' => SmallTuple,
82
80
// b'z' => ShortAscii,
@@ -86,6 +84,9 @@ mod decl {
86
84
}
87
85
}
88
86
87
+ #[ pyattr( name = "version" ) ]
88
+ const VERSION : u32 = 4 ;
89
+
89
90
fn too_short_error ( vm : & VirtualMachine ) -> PyBaseExceptionRef {
90
91
vm. new_exception_msg (
91
92
vm. ctx . exceptions . eof_error . to_owned ( ) ,
@@ -109,93 +110,118 @@ mod decl {
109
110
110
111
/// Dumping helper function to turn a value into bytes.
111
112
fn dump_obj ( buf : & mut Vec < u8 > , value : PyObjectRef , vm : & VirtualMachine ) -> PyResult < ( ) > {
112
- match_class ! ( match value {
113
- pyint @ PyInt => {
114
- if pyint. class( ) . is( vm. ctx. types. bool_type) {
115
- let typ = if pyint. as_bigint( ) . is_zero( ) {
116
- Type :: False
113
+ if vm. is_none ( & value) {
114
+ buf. push ( Type :: None as u8 ) ;
115
+ } else if value. is ( & vm. ctx . ellipsis ) {
116
+ buf. push ( Type :: Ellipsis as u8 ) ;
117
+ } else {
118
+ match_class ! ( match value {
119
+ pyint @ PyInt => {
120
+ if pyint. class( ) . is( vm. ctx. types. bool_type) {
121
+ let typ = if pyint. as_bigint( ) . is_zero( ) {
122
+ Type :: False
123
+ } else {
124
+ Type :: True
125
+ } ;
126
+ buf. push( typ as u8 ) ;
117
127
} else {
118
- Type :: True
119
- } ;
120
- buf. push( typ as u8 ) ;
121
- } else {
122
- buf. push( Type :: Int as u8 ) ;
123
- let ( sign, int_bytes) = pyint. as_bigint( ) . to_bytes_le( ) ;
124
- let mut len = int_bytes. len( ) as i32 ;
125
- if sign == Sign :: Minus {
126
- len = -len;
128
+ buf. push( Type :: Int as u8 ) ;
129
+ let ( sign, int_bytes) = pyint. as_bigint( ) . to_bytes_le( ) ;
130
+ let mut len = int_bytes. len( ) as i32 ;
131
+ if sign == Sign :: Minus {
132
+ len = -len;
133
+ }
134
+ buf. extend( len. to_le_bytes( ) ) ;
135
+ buf. extend( int_bytes) ;
127
136
}
128
- buf. extend( len. to_le_bytes( ) ) ;
129
- buf. extend( int_bytes) ;
130
137
}
131
- }
132
- pyfloat @ PyFloat => {
133
- buf. push( Type :: Float as u8 ) ;
134
- buf. extend( pyfloat. to_f64( ) . to_le_bytes( ) ) ;
135
- }
136
- pystr @ PyStr => {
137
- buf. push( Type :: Str as u8 ) ;
138
- write_size( buf, pystr. as_str( ) . len( ) , vm) ?;
139
- buf. extend( pystr. as_str( ) . as_bytes( ) ) ;
140
- }
141
- pylist @ PyList => {
142
- buf. push( Type :: List as u8 ) ;
143
- let pylist_items = pylist. borrow_vec( ) ;
144
- dump_seq( buf, pylist_items. iter( ) , vm) ?;
145
- }
146
- pyset @ PySet => {
147
- buf. push( Type :: Set as u8 ) ;
148
- let elements = pyset. elements( ) ;
149
- dump_seq( buf, elements. iter( ) , vm) ?;
150
- }
151
- pyfrozen @ PyFrozenSet => {
152
- buf. push( Type :: FrozenSet as u8 ) ;
153
- let elements = pyfrozen. elements( ) ;
154
- dump_seq( buf, elements. iter( ) , vm) ?;
155
- }
156
- pytuple @ PyTuple => {
157
- buf. push( Type :: Tuple as u8 ) ;
158
- dump_seq( buf, pytuple. iter( ) , vm) ?;
159
- }
160
- pydict @ PyDict => {
161
- buf. push( Type :: Dict as u8 ) ;
162
- write_size( buf, pydict. len( ) , vm) ?;
163
- for ( key, value) in pydict {
164
- dump_obj( buf, key, vm) ?;
165
- dump_obj( buf, value, vm) ?;
138
+ pyfloat @ PyFloat => {
139
+ buf. push( Type :: Float as u8 ) ;
140
+ buf. extend( pyfloat. to_f64( ) . to_le_bytes( ) ) ;
166
141
}
167
- }
168
- bytes @ PyByteArray => {
169
- buf. push( Type :: Bytes as u8 ) ;
170
- let data = bytes. borrow_buf( ) ;
171
- write_size( buf, data. len( ) , vm) ?;
172
- buf. extend( & * data) ;
173
- }
174
- co @ PyCode => {
175
- buf. push( Type :: Code as u8 ) ;
176
- let bytes = co. code. map_clone_bag( & bytecode:: BasicBag ) . to_bytes( ) ;
177
- write_size( buf, bytes. len( ) , vm) ?;
178
- buf. extend( bytes) ;
179
- }
180
- _ => {
181
- return Err ( vm. new_not_implemented_error(
182
- "TODO: not implemented yet or marshal unsupported type" . to_owned( ) ,
183
- ) ) ;
184
- }
185
- } ) ;
142
+ pystr @ PyStr => {
143
+ buf. push( if pystr. is_ascii( ) {
144
+ Type :: Ascii
145
+ } else {
146
+ Type :: Unicode
147
+ } as u8 ) ;
148
+ write_size( buf, pystr. as_str( ) . len( ) , vm) ?;
149
+ buf. extend( pystr. as_str( ) . as_bytes( ) ) ;
150
+ }
151
+ pylist @ PyList => {
152
+ buf. push( Type :: List as u8 ) ;
153
+ let pylist_items = pylist. borrow_vec( ) ;
154
+ dump_seq( buf, pylist_items. iter( ) , vm) ?;
155
+ }
156
+ pyset @ PySet => {
157
+ buf. push( Type :: Set as u8 ) ;
158
+ let elements = pyset. elements( ) ;
159
+ dump_seq( buf, elements. iter( ) , vm) ?;
160
+ }
161
+ pyfrozen @ PyFrozenSet => {
162
+ buf. push( Type :: FrozenSet as u8 ) ;
163
+ let elements = pyfrozen. elements( ) ;
164
+ dump_seq( buf, elements. iter( ) , vm) ?;
165
+ }
166
+ pytuple @ PyTuple => {
167
+ buf. push( Type :: Tuple as u8 ) ;
168
+ dump_seq( buf, pytuple. iter( ) , vm) ?;
169
+ }
170
+ pydict @ PyDict => {
171
+ buf. push( Type :: Dict as u8 ) ;
172
+ write_size( buf, pydict. len( ) , vm) ?;
173
+ for ( key, value) in pydict {
174
+ dump_obj( buf, key, vm) ?;
175
+ dump_obj( buf, value, vm) ?;
176
+ }
177
+ }
178
+ bytes @ PyBytes => {
179
+ buf. push( Type :: Bytes as u8 ) ;
180
+ let data = bytes. as_bytes( ) ;
181
+ write_size( buf, data. len( ) , vm) ?;
182
+ buf. extend( & * data) ;
183
+ }
184
+ bytes @ PyByteArray => {
185
+ buf. push( Type :: Bytes as u8 ) ;
186
+ let data = bytes. borrow_buf( ) ;
187
+ write_size( buf, data. len( ) , vm) ?;
188
+ buf. extend( & * data) ;
189
+ }
190
+ co @ PyCode => {
191
+ buf. push( Type :: Code as u8 ) ;
192
+ let bytes = co. code. map_clone_bag( & bytecode:: BasicBag ) . to_bytes( ) ;
193
+ write_size( buf, bytes. len( ) , vm) ?;
194
+ buf. extend( bytes) ;
195
+ }
196
+ _ => {
197
+ return Err ( vm. new_not_implemented_error(
198
+ "TODO: not implemented yet or marshal unsupported type" . to_owned( ) ,
199
+ ) ) ;
200
+ }
201
+ } )
202
+ }
186
203
Ok ( ( ) )
187
204
}
188
205
189
206
#[ pyfunction]
190
- fn dumps ( value : PyObjectRef , vm : & VirtualMachine ) -> PyResult < PyBytes > {
207
+ fn dumps (
208
+ value : PyObjectRef ,
209
+ _version : OptionalArg < i32 > ,
210
+ vm : & VirtualMachine ,
211
+ ) -> PyResult < PyBytes > {
191
212
let mut buf = Vec :: new ( ) ;
192
213
dump_obj ( & mut buf, value, vm) ?;
193
214
Ok ( PyBytes :: from ( buf) )
194
215
}
195
216
196
217
#[ pyfunction]
197
- fn dump ( value : PyObjectRef , f : PyObjectRef , vm : & VirtualMachine ) -> PyResult < ( ) > {
198
- let dumped = dumps ( value, vm) ?;
218
+ fn dump (
219
+ value : PyObjectRef ,
220
+ f : PyObjectRef ,
221
+ version : OptionalArg < i32 > ,
222
+ vm : & VirtualMachine ,
223
+ ) -> PyResult < ( ) > {
224
+ let dumped = dumps ( value, version, vm) ?;
199
225
vm. call_method ( & f, "write" , ( dumped, ) ) ?;
200
226
Ok ( ( ) )
201
227
}
@@ -248,8 +274,10 @@ mod decl {
248
274
let typ = Type :: try_from ( * type_indicator)
249
275
. map_err ( |_| vm. new_value_error ( "bad marshal data (unknown type code)" . to_owned ( ) ) ) ?;
250
276
let ( obj, buf) = match typ {
251
- Type :: True => ( ( true ) . to_pyobject ( vm) , buf) ,
252
- Type :: False => ( ( false ) . to_pyobject ( vm) , buf) ,
277
+ Type :: True => ( true . to_pyobject ( vm) , buf) ,
278
+ Type :: False => ( false . to_pyobject ( vm) , buf) ,
279
+ Type :: None => ( vm. ctx . none ( ) , buf) ,
280
+ Type :: Ellipsis => ( vm. ctx . ellipsis ( ) , buf) ,
253
281
Type :: Int => {
254
282
if buf. len ( ) < 4 {
255
283
return Err ( too_short_error ( vm) ) ;
@@ -276,7 +304,17 @@ mod decl {
276
304
let number = f64:: from_le_bytes ( bytes. try_into ( ) . unwrap ( ) ) ;
277
305
( vm. ctx . new_float ( number) . into ( ) , buf)
278
306
}
279
- Type :: Str => {
307
+ Type :: Ascii => {
308
+ let ( len, buf) = read_size ( buf, vm) ?;
309
+ if buf. len ( ) < len {
310
+ return Err ( too_short_error ( vm) ) ;
311
+ }
312
+ let ( bytes, buf) = buf. split_at ( len) ;
313
+ let s = String :: from_utf8 ( bytes. to_vec ( ) )
314
+ . map_err ( |_| vm. new_value_error ( "invalid utf8 data" . to_owned ( ) ) ) ?;
315
+ ( s. to_pyobject ( vm) , buf)
316
+ }
317
+ Type :: Unicode => {
280
318
let ( len, buf) = read_size ( buf, vm) ?;
281
319
if buf. len ( ) < len {
282
320
return Err ( too_short_error ( vm) ) ;
0 commit comments