Add PyMappingProtocol impl for PyEmbeddings.

Author: sebpuetz

src/embeddings.rs

@@ -10,9 +10,9 @@ use finalfusion::similarity::*;
 use ndarray::Array2;
 use numpy::{IntoPyArray, PyArray1, PyArray2};
 use pyo3::class::iter::PyIterProtocol;
-use pyo3::exceptions;
 use pyo3::prelude::*;
 use pyo3::types::PyTuple;
+use pyo3::{exceptions, PyMappingProtocol};
 use toml::{self, Value};
 
 use crate::{
@@ -114,13 +114,7 @@ impl PyEmbeddings {
     fn embedding(&self, word: &str) -> Option<Py<PyArray1<f32>>> {
         let embeddings = self.embeddings.borrow();
 
-        use EmbeddingsWrap::*;
-        let embedding = match &*embeddings {
-            View(e) => e.embedding(word),
-            NonView(e) => e.embedding(word),
-        };
-
-        embedding.map(|e| {
+        embeddings.embedding(word).map(|e| {
             let gil = pyo3::Python::acquire_gil();
             e.into_owned().into_pyarray(gil.python()).to_owned()
         })
@@ -269,6 +263,21 @@ impl PyEmbeddings {
     }
 }
 
+#[pyproto]
+impl PyMappingProtocol for PyEmbeddings {
+    fn __getitem__(&self, word: &str) -> PyResult<Py<PyArray1<f32>>> {
+        let embeddings = self.embeddings.borrow();
+
+        match embeddings.embedding(word) {
+            Some(embedding) => {
+                let gil = pyo3::Python::acquire_gil();
+                Ok(embedding.into_owned().into_pyarray(gil.python()).to_owned())
+            }
+            None => Err(exceptions::KeyError::py_err("Unknown word and n-grams")),
+        }
+    }
+}
+
 #[pyproto]
 impl PyIterProtocol for PyEmbeddings {
     fn __iter__(slf: PyRefMut<Self>) -> PyResult<PyObject> {

src/embeddings_wrap.rs

@@ -1,5 +1,6 @@
 use finalfusion::norms::NdNorms;
 use finalfusion::prelude::*;
+use finalfusion::storage::CowArray1;
 
 pub enum EmbeddingsWrap {
     NonView(Embeddings<VocabWrap, StorageWrap>),
@@ -30,4 +31,12 @@ impl EmbeddingsWrap {
             View(e) => e.norms(),
         }
     }
+
+    pub fn embedding(&self, word: &str) -> Option<CowArray1<f32>> {
+        use EmbeddingsWrap::*;
+        match self {
+            View(e) => e.embedding(word),
+            NonView(e) => e.embedding(word),
+        }
+    }
 }

tests/test_embedding.py

@@ -28,6 +28,12 @@ def test_embeddings_with_norms_oov(embeddings_fifu):
         "Something out of vocabulary") is None
 
 
+def test_indexing(embeddings_fifu):
+    assert embeddings_fifu["one"] is not None
+    with pytest.raises(KeyError):
+        embeddings_fifu["Something out of vocabulary"]
+
+
 def test_embeddings(embeddings_fifu, embeddings_text):
     for embedding_with_norm, norm in zip(embeddings_fifu, TEST_NORMS):
         unnormed_embed = embedding_with_norm[1] * norm