Properly implement PyProtocols for PyVocab.

In order to get arbitrary keys, PyMappingProtocol::__getitem__
needs to be implemented. To get O(1) __contains__,
PySequenceProtocol::__contains__ needs to be implemented. To get
proper Iteration support, PyIterProtocol::__iter__ needs to be
implemented.

PyO3/pyo3#611

This commit adds the correct implementation of the three traits
to PyVocab.

Author: sebpuetz

src/iter.rs

@@ -80,3 +80,37 @@ impl PyEmbedding {
         self.norm
     }
 }
+
+#[pyclass(name=VocabIterator)]
+pub struct PyVocabIterator {
+    embeddings: Rc<RefCell<EmbeddingsWrap>>,
+    idx: usize,
+}
+
+impl PyVocabIterator {
+    pub fn new(embeddings: Rc<RefCell<EmbeddingsWrap>>, idx: usize) -> Self {
+        PyVocabIterator { embeddings, idx }
+    }
+}
+
+#[pyproto]
+impl PyIterProtocol for PyVocabIterator {
+    fn __iter__(slf: PyRefMut<Self>) -> PyResult<Py<PyVocabIterator>> {
+        Ok(slf.into())
+    }
+
+    fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<String>> {
+        let slf = &mut *slf;
+
+        let embeddings = slf.embeddings.borrow();
+        let vocab = embeddings.vocab();
+
+        if slf.idx < vocab.words_len() {
+            let word = vocab.words()[slf.idx].to_string();
+            slf.idx += 1;
+            Ok(Some(word))
+        } else {
+            Ok(None)
+        }
+    }
+}

src/vocab.rs

@@ -4,9 +4,11 @@ use std::rc::Rc;
 use finalfusion::chunks::vocab::{NGramIndices, SubwordIndices, VocabWrap, WordIndex};
 use finalfusion::prelude::*;
 use pyo3::class::sequence::PySequenceProtocol;
-use pyo3::exceptions;
+use pyo3::exceptions::{IndexError, KeyError, ValueError};
 use pyo3::prelude::*;
+use pyo3::{PyIterProtocol, PyMappingProtocol};
 
+use crate::iter::PyVocabIterator;
 use crate::EmbeddingsWrap;
 
 type NGramIndex = (String, Option<usize>);
@@ -25,16 +27,18 @@ impl PyVocab {
 
 #[pymethods]
 impl PyVocab {
-    fn item_to_indices(&self, key: String) -> Option<PyObject> {
+    #[args(default = "Python::acquire_gil().python().None()")]
+    fn get(&self, key: &str, default: PyObject) -> Option<PyObject> {
         let embeds = self.embeddings.borrow();
-
-        embeds.vocab().idx(key.as_str()).map(|idx| {
-            let gil = pyo3::Python::acquire_gil();
-            match idx {
-                WordIndex::Word(idx) => [idx].to_object(gil.python()),
-                WordIndex::Subword(indices) => indices.to_object(gil.python()),
-            }
-        })
+        let gil = pyo3::Python::acquire_gil();
+        let idx = embeds.vocab().idx(key).map(|idx| match idx {
+            WordIndex::Word(idx) => idx.to_object(gil.python()),
+            WordIndex::Subword(indices) => indices.to_object(gil.python()),
+        });
+        if !default.is_none() && idx.is_none() {
+            return Some(default);
+        }
+        idx
     }
 
     fn ngram_indices(&self, word: &str) -> PyResult<Option<Vec<NGramIndex>>> {
@@ -44,7 +48,7 @@ impl PyVocab {
             VocabWrap::FinalfusionSubwordVocab(inner) => inner.ngram_indices(word),
             VocabWrap::FinalfusionNGramVocab(inner) => inner.ngram_indices(word),
             VocabWrap::SimpleVocab(_) => {
-                return Err(exceptions::ValueError::py_err(
+                return Err(ValueError::py_err(
                     "querying n-gram indices is not supported for this vocabulary",
                 ))
             }
@@ -57,29 +61,71 @@ impl PyVocab {
             VocabWrap::FastTextSubwordVocab(inner) => Ok(inner.subword_indices(word)),
             VocabWrap::FinalfusionSubwordVocab(inner) => Ok(inner.subword_indices(word)),
             VocabWrap::FinalfusionNGramVocab(inner) => Ok(inner.subword_indices(word)),
-            VocabWrap::SimpleVocab(_) => Err(exceptions::ValueError::py_err(
+            VocabWrap::SimpleVocab(_) => Err(ValueError::py_err(
                 "querying subwords' indices is not supported for this vocabulary",
             )),
         }
     }
 }
 
-#[pyproto]
-impl PySequenceProtocol for PyVocab {
-    fn __len__(&self) -> PyResult<usize> {
+impl PyVocab {
+    fn str_to_indices(&self, query: &str) -> PyResult<WordIndex> {
         let embeds = self.embeddings.borrow();
-        Ok(embeds.vocab().words_len())
+        embeds
+            .vocab()
+            .idx(query)
+            .ok_or_else(|| KeyError::py_err(format!("key not found: '{}'", query)))
     }
 
-    fn __getitem__(&self, idx: isize) -> PyResult<String> {
+    fn validate_and_convert_isize_idx(&self, mut idx: isize) -> PyResult<usize> {
         let embeds = self.embeddings.borrow();
-        let words = embeds.vocab().words();
+        let vocab = embeds.vocab();
+        if idx < 0 {
+            idx += vocab.words_len() as isize;
+        }
 
-        if idx >= words.len() as isize || idx < 0 {
-            Err(exceptions::IndexError::py_err("list index out of range"))
+        if idx >= vocab.words_len() as isize || idx < 0 {
+            Err(IndexError::py_err("list index out of range"))
         } else {
-            Ok(words[idx as usize].clone())
+            Ok(idx as usize)
+        }
+    }
+}
+
+#[pyproto]
+impl PyMappingProtocol for PyVocab {
+    fn __getitem__(&self, query: PyObject) -> PyResult<PyObject> {
+        let embeds = self.embeddings.borrow();
+        let vocab = embeds.vocab();
+        let gil = Python::acquire_gil();
+        if let Ok(idx) = query.extract::<isize>(gil.python()) {
+            let idx = self.validate_and_convert_isize_idx(idx)?;
+            return Ok(vocab.words()[idx].clone().into_py(gil.python()));
+        }
+
+        if let Ok(query) = query.extract::<String>(gil.python()) {
+            return self.str_to_indices(&query).map(|idx| match idx {
+                WordIndex::Subword(indices) => indices.into_py(gil.python()),
+                WordIndex::Word(idx) => idx.into_py(gil.python()),
+            });
         }
+
+        Err(KeyError::py_err("key must be integer or string"))
+    }
+}
+
+#[pyproto]
+impl PyIterProtocol for PyVocab {
+    fn __iter__(slf: PyRefMut<Self>) -> PyResult<PyVocabIterator> {
+        Ok(PyVocabIterator::new(slf.embeddings.clone(), 0))
+    }
+}
+
+#[pyproto]
+impl PySequenceProtocol for PyVocab {
+    fn __len__(&self) -> PyResult<usize> {
+        let embeds = self.embeddings.borrow();
+        Ok(embeds.vocab().words_len())
     }
 
     fn __contains__(&self, word: String) -> PyResult<bool> {

tests/test_vocab.py

@@ -1,3 +1,5 @@
+import pytest
+
 TEST_NGRAM_INDICES = [
     ('tüb',
      14),
@@ -53,9 +55,19 @@
      1007)]
 
 
-def test_embeddings_with_norms_oov(embeddings_fifu):
+def test_get(embeddings_text_dims):
+    vocab = embeddings_text_dims.vocab()
+    assert vocab.get("one") is 0
+
+
+def test_get_oov(embeddings_fifu):
+    vocab = embeddings_fifu.vocab()
+    assert vocab.get("Something out of vocabulary") is None
+
+
+def test_get_oov_with_default(embeddings_fifu):
     vocab = embeddings_fifu.vocab()
-    assert vocab.item_to_indices("Something out of vocabulary") is None
+    assert vocab.get("Something out of vocabulary", default=-1) == -1
 
 
 def test_ngram_indices(subword_fifu):
@@ -72,3 +84,41 @@ def test_subword_indices(subword_fifu):
     for subword_index, test_ngram_index in zip(
             subword_indices, TEST_NGRAM_INDICES):
         assert subword_index == test_ngram_index[1]
+
+
+def test_int_idx(embeddings_text_dims):
+    vocab = embeddings_text_dims.vocab()
+    assert vocab[0] == "one"
+
+
+def test_int_idx_out_of_range(embeddings_text_dims):
+    vocab = embeddings_text_dims.vocab()
+    with pytest.raises(IndexError):
+        _ = vocab[42]
+
+
+def test_negative_int_idx(embeddings_text_dims):
+    vocab = embeddings_text_dims.vocab()
+    assert vocab[-1] == "seven"
+
+
+def test_negative_int_idx_out_of_range(embeddings_text_dims):
+    vocab = embeddings_text_dims.vocab()
+    with pytest.raises(IndexError):
+        _ = vocab[-42]
+
+
+def test_string_idx(embeddings_text_dims):
+    vocab = embeddings_text_dims.vocab()
+    assert vocab["one"] == 0
+
+
+def test_string_oov(embeddings_text_dims):
+    vocab = embeddings_text_dims.vocab()
+    with pytest.raises(KeyError):
+        vocab["definitely in vocab"]
+
+
+def test_string_oov_subwords(subword_fifu):
+    vocab = subword_fifu.vocab()
+    assert sorted(vocab["tübingen"]) == [x[1] for x in TEST_NGRAM_INDICES]