Add support for unmarshalling ascii strings.

Author: progval

src/marshal.rs

@@ -30,7 +30,7 @@ pub enum Object {
     //BinaryFloat,
     //Complex,
     //BinaryComplex,
-    //String,
+    String(::std::string::String),
     //Interned,
     //Ref_,
     //Tuple,
@@ -43,7 +43,7 @@ pub enum Object {
     //FrozenSet,
     //Ref,
 
-    //Ascii,
+    Bytes(Vec<u8>), // aka. ASCII in CPython's marshal
     //AsciiInterned,
     //SmallTuple,
     //ShortAscii,
@@ -66,6 +66,29 @@ macro_rules! read_byte {
     }};
 }
 
+fn read_long<R: io::Read>(r: &mut R) -> Result<u32, UnmarshalError> {
+    let mut buf = [0, 0, 0, 0];
+    match r.read_exact(&mut buf) {
+        Err(err) => return Err(UnmarshalError::Io(err)),
+        Ok(()) => Ok(buf[0] as u32 + 256*(buf[1] as u32 + 256*(buf[2] as u32 + 256*(buf[3] as u32))))
+    }
+}
+
+fn read_ascii_string<R: io::Read>(r: &mut R, size: usize) -> Result<String, UnmarshalError> {
+    let mut buf = Vec::<u8>::new();
+    buf.resize(size, 0);
+    match r.read_exact(&mut buf) {
+        Err(err) => return Err(UnmarshalError::Io(err)),
+        Ok(()) => ()
+    };
+    let mut string = String::new();
+    string.reserve(buf.len()); // The string will use more bytes than this if there are extended ascii characters, but it can't hurt to reserve anyway
+    for c in buf {
+        string.push(c as char);
+    }
+    Ok(string)
+}
+
 pub fn read_object<R: io::Read>(r: &mut R) -> Result<Object, UnmarshalError> {
     let byte = read_byte!(r);
     let _flag = byte & 0b10000000; // TODO: do something with this
@@ -75,13 +98,10 @@ pub fn read_object<R: io::Read>(r: &mut R) -> Result<Object, UnmarshalError> {
         'N' => Object::None,
         'F' => Object::False,
         'T' => Object::True,
-        'i' => {
-            let mut buf = [0, 0, 0, 0];
-            match r.read_exact(&mut buf) {
-                Err(err) => return Err(UnmarshalError::Io(err)),
-                Ok(()) => ()
-            };
-            Object::Int(buf[0] as u32 + 256*(buf[1] as u32 + 256*(buf[2] as u32 + 256*(buf[3] as u32))))
+        'i' => Object::Int(try!(read_long(r))),
+        'z' | 'Z' => { // “short ascii”, “short ascii interned”
+            let size = read_byte!(r) as usize;
+            Object::String(try!(read_ascii_string(r, size)))
         },
         _ => panic!(format!("Unsupported opcode: {}", opcode as char)),
     };
@@ -111,3 +131,12 @@ fn test_int() {
     let mut reader: &[u8] = b"\xe9\xe8\x03\x00\x00";
     assert_eq!(Object::Int(1000), read_object(&mut reader).unwrap());
 }
+
+#[test]
+fn test_string() {
+    let mut reader: &[u8] = b"\xda\x03foo";
+    assert_eq!(Object::String("foo".to_string()), read_object(&mut reader).unwrap());
+
+    let mut reader: &[u8] = b"\xda\x04foo\xe9"; // Note: this string was not generated with the marshal module
+    assert_eq!(Object::String("fooé".to_string()), read_object(&mut reader).unwrap());
+}