add utility methods to the bencode objects

This commit is contained in:
2016-12-11 19:15:09 -05:00
parent 0fc6704cc2
commit bd021f80a3
3 changed files with 175 additions and 32 deletions

View File

@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
use std::num::ParseIntError;
use std::str::{self, Utf8Error};
use bencode::Object;
use bencode::{Bytes, Dict, List, Object};
use buffer::Buffer;
#[derive(Debug)]
@@ -46,7 +46,7 @@ fn decode_object(buf: &mut Buffer) -> DecodeResult<Object> {
}
list.push(decode_object(buf)?);
}
Ok(Object::List(list))
Ok(List::new(list))
}
Some(b'd') => {
buf.advance(1);
@@ -58,9 +58,9 @@ fn decode_object(buf: &mut Buffer) -> DecodeResult<Object> {
}
let key = _decode_bytes(buf)?;
let val = decode_object(buf)?;
dict.insert(key, val);
dict.insert(Bytes(key), val);
}
Ok(Object::Dict(dict))
Ok(Dict::new(dict))
}
_ => Err(DecodeError),
}
@@ -84,7 +84,7 @@ fn _decode_int(buf: &mut Buffer, term: u8) -> DecodeResult<i64> {
}
fn decode_bytes(buf: &mut Buffer) -> DecodeResult<Object> {
_decode_bytes(buf).map(|bytes| Object::Bytes(bytes))
_decode_bytes(buf).map(|bytes| Bytes::new(bytes))
}
fn _decode_bytes(buf: &mut Buffer) -> Result<Vec<u8>, DecodeError> {
@@ -116,7 +116,7 @@ fn test_int_neg() {
fn test_bytes() {
let mut buf = Buffer::new(b"5:hello");
assert_eq!(decode_bytes(&mut buf).unwrap(), Object::Bytes(b"hello".to_vec()));
assert_eq!(decode_bytes(&mut buf).unwrap(), Bytes::new(b"hello".to_vec()));
assert_eq!(buf.pos(), 7);
}
@@ -127,7 +127,7 @@ fn test_list() {
let obj = decode_object(&mut buf).unwrap();
let list = obj.as_list().unwrap();
assert_eq!(list, vec![Object::Int(1), Object::Int(2), Object::Int(3)]);
assert_eq!(list, &List(vec![Object::Int(1), Object::Int(2), Object::Int(3)]));
assert_eq!(buf.pos(), 11);
}
@@ -138,6 +138,6 @@ fn test_dict() {
let obj = decode_object(&mut buf).unwrap();
let dict = obj.as_dict().unwrap();
assert_eq!(dict[&b"hello"[..]], Object::Int(1337));
assert_eq!(dict.get_int("hello").unwrap(), 1337);
assert_eq!(buf.pos(), 15);
}

View File

@@ -1,7 +1,6 @@
use std::borrow::Borrow;
use std::collections::BTreeMap;
use bencode::Object;
use bencode::{Dict, Object};
pub fn encode<B>(obj: B) -> Vec<u8> where B: Borrow<Object> {
let mut buff = Vec::new();
@@ -35,9 +34,9 @@ fn encode_list(buff: &mut Vec<u8>, list: &[Object]) {
buff.push(b'e');
}
fn encode_dict(buff: &mut Vec<u8>, dict: &BTreeMap<Vec<u8>, Object>) {
fn encode_dict(buff: &mut Vec<u8>, dict: &Dict) {
buff.push(b'd');
for (key, val) in dict {
for (key, val) in dict.iter() {
encode_bytes(buff, key);
encode_object(buff, val);
}
@@ -56,18 +55,25 @@ fn test_int_neg() {
#[test]
fn test_bytes() {
assert_eq!(encode(Object::Bytes(b"hello".to_vec())), b"5:hello");
use bencode::Bytes;
assert_eq!(encode(Bytes::new(b"hello".to_vec())), b"5:hello");
}
#[test]
fn test_list() {
use bencode::List;
let list = vec![Object::Int(1), Object::Int(2), Object::Int(3)];
assert_eq!(encode(Object::List(list)), b"li1ei2ei3ee");
assert_eq!(encode(List::new(list)), b"li1ei2ei3ee");
}
#[test]
fn test_dict() {
use std::collections::BTreeMap;
use bencode::{Bytes, Dict};
let mut dict = BTreeMap::new();
dict.insert(b"hello".to_vec(), Object::Int(1337));
assert_eq!(encode(Object::Dict(dict)), b"d5:helloi1337ee")
dict.insert(Bytes(b"hello".to_vec()), Object::Int(1337));
assert_eq!(encode(Dict::new(dict)), b"d5:helloi1337ee")
}

View File

@@ -1,46 +1,56 @@
mod decode;
mod encode;
use std::borrow::Cow;
use std::collections::BTreeMap;
use std::fmt;
use std::ops::{Deref, DerefMut};
use std::str;
pub use self::decode::{decode, DecodeError, DecodeResult};
pub use self::encode::encode;
#[derive(Eq, PartialEq)]
#[derive(Clone, Eq, PartialEq)]
pub enum Object {
Int(i64),
Bytes(Vec<u8>),
List(Vec<Object>),
Dict(BTreeMap<Vec<u8>, Object>),
Bytes(Bytes),
List(List),
Dict(Dict),
}
impl Object {
pub fn as_int(self) -> Option<i64> {
match self {
pub fn as_int(&self) -> Option<i64> {
match *self {
Object::Int(num) => Some(num),
_ => None,
}
}
pub fn as_bytes(self) -> Option<Vec<u8>> {
match self {
Object::Bytes(bytes) => Some(bytes),
pub fn as_bytes(&self) -> Option<&Bytes> {
match *self {
Object::Bytes(ref bytes) => Some(bytes),
_ => None,
}
}
pub fn as_list(self) -> Option<Vec<Object>> {
match self {
Object::List(list) => Some(list),
pub fn as_str(&self) -> Option<&str> {
self.as_bytes().and_then(|b| str::from_utf8(b).ok())
}
pub fn as_string(&self) -> Option<String> {
self.as_str().map(|s| s.to_string())
}
pub fn as_list(&self) -> Option<&List> {
match *self {
Object::List(ref list) => Some(list),
_ => None,
}
}
pub fn as_dict(self) -> Option<BTreeMap<Vec<u8>, Object>> {
match self {
Object::Dict(dict) => Some(dict),
pub fn as_dict(&self) -> Option<&Dict> {
match *self {
Object::Dict(ref dict) => Some(dict),
_ => None,
}
}
@@ -66,9 +76,136 @@ impl fmt::Debug for Object {
}
Object::Dict(ref dict) => {
f.debug_map()
.entries(dict.iter().map(|(k, v)| (Object::Bytes(k.clone()), v)))
.entries(dict.iter().map(|(k, v)| (k.clone(), v)))
.finish()
}
}
}
}
#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct Bytes(Vec<u8>);
impl Bytes {
pub fn new(bytes: Vec<u8>) -> Object {
Object::Bytes(Bytes(bytes))
}
pub fn str(&self) -> Option<&str> {
str::from_utf8(&self.0).ok()
}
pub fn string(&self) -> Option<String> {
self.str().map(|s| s.to_string())
}
}
impl Deref for Bytes {
type Target = Vec<u8>;
fn deref(&self) -> &Vec<u8> {
&self.0
}
}
impl DerefMut for Bytes {
fn deref_mut(&mut self) -> &mut Vec<u8> {
&mut self.0
}
}
pub trait AsBytes {
fn as_bytes(&self) -> Cow<Bytes>;
}
impl AsBytes for Bytes {
fn as_bytes(&self) -> Cow<Bytes> {
Cow::Borrowed(self)
}
}
impl<'a> AsBytes for &'a str {
fn as_bytes(&self) -> Cow<Bytes> {
Cow::Owned(Bytes(str::as_bytes(self).to_vec()))
}
}
impl<'a> AsBytes for &'a [u8] {
fn as_bytes(&self) -> Cow<Bytes> {
Cow::Owned(Bytes(self.to_vec()))
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct List(Vec<Object>);
impl List {
pub fn new(list: Vec<Object>) -> Object {
Object::List(List(list))
}
}
impl Deref for List {
type Target = Vec<Object>;
fn deref(&self) -> &Vec<Object> {
&self.0
}
}
impl DerefMut for List {
fn deref_mut(&mut self) -> &mut Vec<Object> {
&mut self.0
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Dict(BTreeMap<Bytes, Object>);
impl Dict {
pub fn new(dict: BTreeMap<Bytes, Object>) -> Object {
Object::Dict(Dict(dict))
}
pub fn get_object<A>(&self, key: A) -> Option<&Object> where A: AsBytes {
self.0.get(&key.as_bytes())
}
pub fn get_int<A>(&self, key: A) -> Option<i64> where A: AsBytes {
self.0.get(&key.as_bytes()).and_then(|o| o.as_int())
}
pub fn get_bytes<A>(&self, key: A) -> Option<&Bytes> where A: AsBytes {
self.0.get(&key.as_bytes()).and_then(|o| o.as_bytes())
}
pub fn get_str<A>(&self, key: A) -> Option<&str> where A: AsBytes {
self.0.get(&key.as_bytes()).and_then(|o| o.as_bytes().and_then(|b| str::from_utf8(b).ok()))
}
pub fn get_string<A>(&self, key: A) -> Option<String> where A: AsBytes {
self.get_str(key).map(|s| s.to_string())
}
pub fn get_list<A>(&self, key: A) -> Option<&List> where A: AsBytes {
self.0.get(&key.as_bytes()).and_then(|o| o.as_list())
}
pub fn get_dict<A>(&self, key: A) -> Option<&Dict> where A: AsBytes {
self.0.get(&key.as_bytes()).and_then(|o| o.as_dict())
}
}
impl Deref for Dict {
type Target = BTreeMap<Bytes, Object>;
fn deref(&self) -> &BTreeMap<Bytes, Object> {
&self.0
}
}
impl DerefMut for Dict {
fn deref_mut(&mut self) -> &mut BTreeMap<Bytes, Object> {
&mut self.0
}
}