store uitgewerkt

src/store.rs

@@ -3,24 +3,46 @@ use crate::descriptors::Descriptor;
 use std::fs;
 use std::path::Path;
 use std::fmt;
-use log::info;
+use log::{info, debug, error};
+
+#[derive(Debug)]
+pub enum SaveError {
+	Serialization,
+	File,
+}
 
 pub struct DescriptorStore {
-	map: HashMap<String, u64>
+	map: HashMap<String, u64>,
+	save_location: std::path::PathBuf,
 }
 
 impl DescriptorStore {
-	/// Makes a new DescriptorStore.
-	/// Reads from file if possible, else empty
-	pub fn new() -> DescriptorStore {
-		//TODO: Set directory and file location from global config
-		let count = fs::read_dir("img").unwrap().count();
-		let map_file = fs::read("map.messagepack");
-		let map: HashMap<String, u64> = match map_file {
-			Ok(f) => rmp_serde::from_slice(&f).unwrap(),
-			Err(_e) => HashMap::with_capacity(count),
+	/// Makes a new empty DescriptorStore with default settings
+	pub fn new() -> Self {
+		let map: HashMap<String, u64> = HashMap::new();
+		DescriptorStore { map, save_location: std::path::PathBuf::from("store.messagepack")}
+	}
+
+	/// Sets the file location and loads data from file (if available)
+	pub fn with_file<P: AsRef<Path>>(mut self, path: P) -> Self {
+		self.save_location = std::path::PathBuf::from(path.as_ref());
+		let map_file = fs::read(&self.save_location);
+		match map_file {
+			Ok(f) => self.map = rmp_serde::from_slice(&f).unwrap(),
+			Err(_) => info!("{} not found, starting from empty store.", self.save_location.display()),
 		};
-		DescriptorStore { map }
+		self
+	}
+
+	pub fn save(&self) -> Result<(), SaveError> {
+		let serialized: Vec<u8> = match rmp_serde::to_vec(&self.map) {
+			Ok(value) => value,
+			Err(_e) => return Err(SaveError::Serialization),
+		};
+		match fs::write(&self.save_location, &serialized) {
+			Ok(_) => Ok(()),
+			Err(_e) => Err(SaveError::File),
+		}
 	}
 
 	/// Returns true iff the store already contains the key
@@ -29,22 +51,37 @@ impl DescriptorStore {
 	}
 
 	/// Inserts a single value into the store
-	pub fn insert(mut self, key: String, value: u64) {
+	pub fn insert(&mut self, key: String, value: u64) {
 		self.map.insert(key, value);
 	}
 
 
-	/// Calculates all descriptions with a given descriptor for a folder
+	/// Calculates all descriptions with a given descriptor 
 	pub fn insert_directory<T: Descriptor, U: AsRef<Path>>(&mut self, dir: U, desc: T) {
 		for node in fs::read_dir(dir).unwrap() {
 			let file = node.expect("Error walking directory");
-			let name = file.file_name().into_string().expect("Issue with filename");
+			let name = match file.file_name().into_string() {
+				Ok(v) => v,
+				Err(e) => {
+					error!("Error reading {}: {:?}:", file.path().to_string_lossy(), e);
+					continue
+				}
+			};
 			if !self.contains(name.to_string()) {
 				info!("Processing {}", name);
-				let img = image::open(file.path()).expect("Unable to open file");
+				let img = match image::open(file.path()) {
+					Ok(v) => v,
+					Err(e) => {
+						error!("Failed to process {}: {}", name, e);
+						continue
+					}
+				};
 				let phash = desc.describe(img);
 				self.map.insert(name, phash);
+			} else {
+				debug!("{} already known, skipping.", name);
 			}
+			self.save().expect("error");
 		}
 	}
 }
@@ -53,9 +90,9 @@ impl fmt::Display for DescriptorStore {
 	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 		let mut output = String::new();
 		for (key, val) in self.map.iter() {
-			let entry = format!("{key: >15}: {val}\t {val:064b}\n");
+			let entry = format!("{key:020}: {val}\t {val:064b}\n");
 			output.push_str(&entry);
 		}
-		write!(f, "DescriptorStore:\n{}", output)
+		write!(f, "DescriptorStore:\n{}Total: {}", output, self.map.len())
 	}
 }