From d075eb00fff60603556a29714c348821a0bdda9a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81rni=20Dagur?= <agudmundsson@fc-md.umd.edu>
Date: Sun, 29 Mar 2020 23:23:03 -0400
Subject: Add serialization and deserialization methods

---
 src/lib.rs | 78 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 76 insertions(+), 2 deletions(-)

(limited to 'src/lib.rs')

diff --git a/src/lib.rs b/src/lib.rs
index 86512f7..5ba8d8a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -10,11 +10,17 @@
     unused_qualifications
 )]
 
+use adblock::blocker::BlockerError as RustBlockerError;
+use adblock::blocker::BlockerResult as RustBlockerResult;
+use adblock::engine::Engine as RustEngine;
+use failure::Fail;
 use pyo3::class::PyObjectProtocol;
+use pyo3::exceptions::ValueError as PyValueError;
 use pyo3::prelude::*;
+use pyo3::PyErr;
 
-use adblock::blocker::BlockerResult as RustBlockerResult;
-use adblock::engine::Engine as RustEngine;
+use std::fs;
+use std::io::{Read, Write};
 
 /// Brave's adblocking library in Python!
 #[pymodule]
@@ -72,6 +78,38 @@ impl PyObjectProtocol for BlockerResult {
     }
 }
 
+#[derive(Fail, Debug, PartialEq, Copy, Clone)]
+pub enum BlockerError {
+    #[fail(display = "Serialization error")]
+    SerializationError,
+    #[fail(display = "Deserialization error")]
+    DeserializationError,
+    #[fail(display = "Optimized filter exists")]
+    OptimizedFilterExistence,
+    #[fail(display = "Bad filter add unsupported")]
+    BadFilterAddUnsupported,
+    #[fail(display = "Filter exists")]
+    FilterExists,
+}
+
+impl Into<PyErr> for BlockerError {
+    fn into(self) -> PyErr {
+        PyErr::new::<PyValueError, _>(format!("{:?}", self))
+    }
+}
+
+impl Into<BlockerError> for RustBlockerError {
+    fn into(self) -> BlockerError {
+        match self {
+            Self::SerializationError => BlockerError::SerializationError,
+            Self::DeserializationError => BlockerError::DeserializationError,
+            Self::OptimizedFilterExistence => BlockerError::OptimizedFilterExistence,
+            Self::BadFilterAddUnsupported => BlockerError::BadFilterAddUnsupported,
+            Self::FilterExists => BlockerError::FilterExists,
+        }
+    }
+}
+
 #[pyclass]
 pub struct Engine {
     engine: RustEngine,
@@ -148,6 +186,42 @@ impl Engine {
         blocker_result.into()
     }
 
+    pub fn serialize(&mut self) -> PyResult<Vec<u8>> {
+        let result = self.engine.serialize();
+        match result {
+            Ok(x) => Ok(x),
+            Err(error) => {
+                let my_blocker_error: BlockerError = error.into();
+                Err(my_blocker_error.into())
+            }
+        }
+    }
+
+    pub fn serialize_to_file(&mut self, file: &str) -> PyResult<()> {
+        let mut fd = fs::File::open(file)?;
+        let data = self.serialize()?;
+        fd.write_all(&data)?;
+        Ok(())
+    }
+
+    pub fn deserialize(&mut self, serialized: &[u8]) -> PyResult<()> {
+        let result = self.engine.deserialize(serialized);
+        match result {
+            Ok(x) => Ok(x),
+            Err(error) => {
+                let my_blocker_error: BlockerError = error.into();
+                Err(my_blocker_error.into())
+            }
+        }
+    }
+
+    pub fn deserialize_from_file(&mut self, file: &str) -> PyResult<()> {
+        let mut fd = fs::File::open(file)?;
+        let mut data: Vec<u8> = Vec::new();
+        fd.read_to_end(&mut data)?;
+        self.deserialize(&data)
+    }
+
     pub fn filter_exists(&self, filter: &str) -> bool {
         self.engine.filter_exists(filter)
     }
-- 
cgit v1.2.3