From b9cdf3d85e0ef28cf89a1a307a984879ba604511 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Sat, 24 Dec 2022 12:36:54 +0000
Subject: String cache

---
 rust/src/tree_cache/binding.rs | 64 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 63 insertions(+), 1 deletion(-)

(limited to 'rust/src/tree_cache/binding.rs')

diff --git a/rust/src/tree_cache/binding.rs b/rust/src/tree_cache/binding.rs
index 70207f8781..e01601daf5 100644
--- a/rust/src/tree_cache/binding.rs
+++ b/rust/src/tree_cache/binding.rs
@@ -2,7 +2,9 @@ use std::hash::Hash;
 
 use anyhow::Error;
 use pyo3::{
-    pyclass, pymethods, types::PyModule, IntoPy, PyAny, PyObject, PyResult, Python, ToPyObject,
+    pyclass, pymethods,
+    types::{PyIterator, PyModule},
+    IntoPy, PyAny, PyObject, PyResult, Python, ToPyObject,
 };
 
 use super::TreeCache;
@@ -10,6 +12,7 @@ use super::TreeCache;
 pub fn register_module(py: Python<'_>, m: &PyModule) -> PyResult<()> {
     let child_module = PyModule::new(py, "tree_cache")?;
     child_module.add_class::<PythonTreeCache>()?;
+    child_module.add_class::<StringTreeCache>()?;
 
     m.add_submodule(child_module)?;
 
@@ -126,3 +129,62 @@ impl PythonTreeCache {
         todo!()
     }
 }
+
+#[pyclass]
+struct StringTreeCache(TreeCache<String, String>);
+
+#[pymethods]
+impl StringTreeCache {
+    #[new]
+    fn new() -> Self {
+        StringTreeCache(Default::default())
+    }
+
+    pub fn set(&mut self, key: &PyAny, value: String) -> Result<(), Error> {
+        let key = key
+            .iter()?
+            .map(|o| o.expect("iter failed").extract().expect("not a string"));
+
+        self.0.set(key, value)?;
+
+        Ok(())
+    }
+
+    // pub fn get_node(&self, key: &PyAny) -> Result<Option<&TreeCacheNode<K, PyObject>>, Error> {
+    //     todo!()
+    // }
+
+    pub fn get(&self, key: &PyAny) -> Result<Option<&String>, Error> {
+        let key = key.iter()?.map(|o| {
+            o.expect("iter failed")
+                .extract::<String>()
+                .expect("not a string")
+        });
+
+        Ok(self.0.get(key)?)
+    }
+
+    // pub fn pop_node(&mut self, key: &PyAny) -> Result<Option<TreeCacheNode<K, PyObject>>, Error> {
+    //     todo!()
+    // }
+
+    pub fn pop(&mut self, key: Vec<String>) -> Result<Option<String>, Error> {
+        Ok(self.0.pop(&key)?)
+    }
+
+    pub fn clear(&mut self) {
+        self.0.clear()
+    }
+
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
+
+    pub fn values(&self) -> Vec<&String> {
+        self.0.values().collect()
+    }
+
+    pub fn items(&self) -> Vec<(Vec<&HashablePyObject>, &PyObject)> {
+        todo!()
+    }
+}
-- 
cgit 1.5.1