Repository: WorksApplications/sudachi.rs (language: python)

Issue:
Python API: POS matcher
Implement a way to match a morpheme to a set of POS tags.

Need to decide an interface.

Callback interface which uses a callable to check eligible POS during creation.

```python
matcher = dic.pos_matcher(lambda p: p[0] == "名詞")
```

Declarative interface which expects user to pass a list of POS tuples.

```python
matcher = dic.pos_matcher([('名詞', '一般', '*', '*', '*', '*')])
```

With maybe None (and skipped fields) matching any POS component.
```dic.pos_matcher([('名詞', None)])``` will match any noun.


Discussion / hints:

Fix patch:
diff --git a/python/py_src/sudachipy/sudachipy.pyi b/python/py_src/sudachipy/sudachipy.pyi
index 0e5fd719..82b632c8 100644
--- a/python/py_src/sudachipy/sudachipy.pyi
+++ b/python/py_src/sudachipy/sudachipy.pyi
@@ -1,7 +1,20 @@
-from typing import ClassVar, Iterator, List, Set, Literal, Optional
+from typing import ClassVar, Iterator, List, Tuple, Union, Callable, Iterable, Optional, Literal, Set
 
 import sudachipy
 
+POS = Tuple[str, str, str, str, str, str]
+# POS element
+PE = Optional[str]
+PartialPOS = Union[
+    Tuple[PE, PE, PE, PE, PE, PE],
+    Tuple[PE, PE, PE, PE, PE],
+    Tuple[PE, PE, PE, PE],
+    Tuple[PE, PE, PE],
+    Tuple[PE, PE],
+    Tuple[PE],
+    Tuple[()],
+]
+
 FieldSet = Optional[Set[Literal["surface", "pos", "normalized_form", "dictionary_form", "reading_form",
                                 "word_structure", "split_a", "split_b", "synonym_group_id"]]]
 
@@ -38,6 +51,22 @@ class Dictionary:
         """
         ...
 
+    def pos_matcher(self, target: Union[Iterable[PartialPOS], Callable[[POS], bool]]) -> PosMatcher:
+        """
+        Creates a new POS matcher
+
+        If target is a function, then it must return whether a POS should match or not.
+        If target a list, it should contain partially specified POS.
+        By partially specified it means that it is possible to omit POS fields or
+        use None as a sentinel value that matches any POS.
+
+        For example, ('名詞',) will match any noun and
+        (None, None, None, None, None, '終止形') will match any word in 終止形 conjugation form.
+
+        :param target: can be either a function or a list of POS tuples.
+        """
+        ...
+
 
 class Morpheme:
     """
@@ -88,7 +117,7 @@ class Morpheme:
         """
         ...
 
-    def part_of_speech(self) -> List[str]:
+    def part_of_speech(self) -> POS:
         """
         Returns the part of speech.
         """
@@ -210,3 +239,14 @@ class WordInfo:
     @classmethod
     def __init__(self) -> None: ...
     def length(self) -> int: ...
+
+class PosMatcher:
+    def __iter__(self) -> Iterator[POS]: ...
+    def __len__(self) -> int: ...
+    def __call__(self, m: Morpheme) -> bool:
+        """
+        Checks whether a morpheme has matching POS
+        :param m: morpheme
+        :return: if morpheme has matching POS
+        """
+        ...
\ No newline at end of file
diff --git a/python/src/dictionary.rs b/python/src/dictionary.rs
index 733eb588..56075796 100644
--- a/python/src/dictionary.rs
+++ b/python/src/dictionary.rs
@@ -22,6 +22,8 @@ use pyo3::prelude::*;
 use pyo3::types::{PySet, PyString, PyTuple};
 
 use sudachi::analysis::stateless_tokenizer::DictionaryAccess;
+
+use crate::pos_matcher::PyPosMatcher;
 use sudachi::config::Config;
 use sudachi::dic::dictionary::JapaneseDictionary;
 use sudachi::dic::grammar::Grammar;
@@ -164,6 +166,14 @@ impl PyDictionary {
         Ok(tok)
     }
 
+    /// Creates a POS matcher object
+    ///
+    /// target can be either a callable or list of POS partial tuples
+    #[pyo3(text_signature = "($self, target")]
+    fn pos_matcher<'py>(&'py self, py: Python<'py>, target: &PyAny) -> PyResult<PyPosMatcher> {
+        PyPosMatcher::create(py, self.dictionary.as_ref().unwrap(), target)
+    }
+
     /// Close this dictionary
     #[pyo3(text_signature = "($self)")]
     fn close(&mut self) {
diff --git a/python/src/lib.rs b/python/src/lib.rs
index 47c36292..2ecd7bea 100644
--- a/python/src/lib.rs
+++ b/python/src/lib.rs
@@ -19,6 +19,7 @@ use pyo3::prelude::*;
 mod build;
 mod dictionary;
 mod morpheme;
+mod pos_matcher;
 mod tokenizer;
 mod word_info;
 
diff --git a/python/src/morpheme.rs b/python/src/morpheme.rs
index 58d9f739..cfecdb74 100644
--- a/python/src/morpheme.rs
+++ b/python/src/morpheme.rs
@@ -189,7 +189,7 @@ impl PyMorpheme {
 
     /// Returns the id of the part of speech in the dictionary
     #[pyo3(text_signature = "($self)")]
-    fn part_of_speech_id(&self) -> u16 {
+    pub fn part_of_speech_id(&self) -> u16 {
         self.list.get_word_info(self.index).pos_id()
     }
 
diff --git a/python/src/pos_matcher.rs b/python/src/pos_matcher.rs
new file mode 100644
index 00000000..4266a9a7
--- /dev/null
+++ b/python/src/pos_matcher.rs
@@ -0,0 +1,171 @@
+/*
+ *  Copyright (c) 2021 Works Applications Co., Ltd.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+use crate::dictionary::PyDicData;
+use crate::morpheme::PyMorpheme;
+use pyo3::exceptions::PyException;
+use pyo3::prelude::*;
+use pyo3::types::{PyBool, PyIterator, PyTuple};
+use std::sync::Arc;
+use sudachi::analysis::stateless_tokenizer::DictionaryAccess;
+use sudachi::pos::PosMatcher;
+
+#[pyclass(name = "PosMatcher", module = "sudachipy")]
+pub struct PyPosMatcher {
+    matcher: PosMatcher,
+    dic: Arc<PyDicData>,
+}
+
+impl PyPosMatcher {
+    pub(crate) fn create<'py>(
+        py: Python<'py>,
+        dic: &'py Arc<PyDicData>,
+        data: &'py PyAny,
+    ) -> PyResult<PyPosMatcher> {
+        if data.is_callable() {
+            Self::create_from_fn(dic, data, py)
+        } else {
+            let iter = data.iter()?;
+            Self::create_from_iter(dic, iter)
+        }
+    }
+
+    fn create_from_fn(dic: &Arc<PyDicData>, func: &PyAny, py: Python) -> PyResult<Self> {
+        let mut data = Vec::new();
+        for (pos_id, pos) in dic.pos.iter().enumerate() {
+            let args = PyTuple::new(py, &[pos]);
+            if func.call1(args)?.cast_as::<PyBool>()?.is_true() {
+                data.push(pos_id as u16);
+            }
+        }
+        Ok(Self {
+            matcher: PosMatcher::new(data),
+            dic: dic.clone(),
+        })
+    }
+
+    fn create_from_iter(dic: &Arc<PyDicData>, data: &PyIterator) -> PyResult<Self> {
+        let mut result = Vec::new();
+        for item in data {
+            let item = item?.cast_as::<PyTuple>()?;
+            Self::match_pos_elements(&mut result, dic.as_ref(), item)?;
+        }
+        Ok(Self {
+            matcher: PosMatcher::new(result),
+            dic: dic.clone(),
+        })
+    }
+
+    fn match_pos_elements(data: &mut Vec<u16>, dic: &PyDicData, elem: &PyTuple) -> PyResult<()> {
+        let start_len = data.len();
+
+        let elen = elem.len();
+        for (pos_id, pos) in dic.grammar().pos_list.iter().enumerate() {
+            let check = |idx: usize| -> PyResult<bool> {
+                let x = elem.get_item(idx)?;
+                if x.is_none() {
+                    return Ok(false);
+                }
+                Ok(x.str()?.to_str()? != pos[idx])
+            };
+            if elen > 0 && check(0)? {
+                continue;
+            }
+            if elen > 1 && check(1)? {
+                continue;
+            }
+            if elen > 2 && check(2)? {
+                continue;
+            }
+            if elen > 3 && check(3)? {
+                continue;
+            }
+            if elen > 4 && check(4)? {
+                continue;
+            }
+            if elen > 5 && check(5)? {
+                continue;
+            }
+            data.push(pos_id as u16);
+        }
+
+        if start_len == data.len() {
+            Err(PyException::new_err(format!(
+                "POS {:?} did not match any elements",
+                elem.repr()?
+            )))
+        } else {
+            Ok(())
+        }
+    }
+}
+
+#[pymethods]
+impl PyPosMatcher {
+    pub fn __call__(&self, m: &PyMorpheme) -> bool {
+        let pos_id = m.part_of_speech_id();
+        self.matcher.matches_id(pos_id)
+    }
+
+    pub fn __str__(&self) -> String {
+        format!("PosMatcher<{} pos>", self.matcher.num_entries())
+    }
+
+    pub fn __iter__(&self) -> PyPosIter {
+        PyPosIter::new(self.matcher.entries(), self.dic.clone())
+    }
+
+    pub fn __len__(&self) -> usize {
+        self.matcher.num_entries()
+    }
+}
+
+#[pyclass(name = "PosMatcherIterator", module = "sudachipy")]
+pub struct PyPosIter {
+    data: Vec<u16>,
+    dic: Arc<PyDicData>,
+    index: usize,
+}
+
+impl PyPosIter {
+    fn new(data: impl Iterator<Item = u16>, dic: Arc<PyDicData>) -> Self {
+        let mut result: Vec<u16> = data.collect();
+        result.sort();
+        Self {
+            data: result,
+            dic,
+            index: 0,
+        }
+    }
+}
+
+#[pymethods]
+impl PyPosIter {
+    fn __iter__(slf: &PyCell<Self>) -> &PyCell<Self> {
+        slf
+    }
+
+    fn __next__<'py>(&'py mut self, py: Python<'py>) -> Option<&'py PyTuple> {
+        let idx = self.index;
+        self.index += 1;
+        if idx >= self.data.len() {
+            return None;
+        }
+        let pos_id = self.data[idx];
+        let pos = &self.dic.pos[pos_id as usize];
+        Some(pos.as_ref(py))
+    }
+}
diff --git a/python/src/tokenizer.rs b/python/src/tokenizer.rs
index bc9d9975..baeff995 100644
--- a/python/src/tokenizer.rs
+++ b/python/src/tokenizer.rs
@@ -21,7 +21,6 @@ use pyo3::prelude::*;
 
 use sudachi::analysis::stateful_tokenizer::StatefulTokenizer;
 
-use sudachi::dic::dictionary::JapaneseDictionary;
 use sudachi::dic::subset::InfoSubset;
 use sudachi::prelude::*;
 
diff --git a/sudachi/src/lib.rs b/sudachi/src/lib.rs
index a5d48657..d4c92696 100644
--- a/sudachi/src/lib.rs
+++ b/sudachi/src/lib.rs
@@ -34,6 +34,7 @@ pub mod sentence_splitter;
 pub(crate) mod util;
 
 mod hash;
+pub mod pos;
 #[cfg(test)]
 pub mod test;
 
diff --git a/sudachi/src/pos.rs b/sudachi/src/pos.rs
new file mode 100644
index 00000000..2136a485
--- /dev/null
+++ b/sudachi/src/pos.rs
@@ -0,0 +1,46 @@
+/*
+ *  Copyright (c) 2021 Works Applications Co., Ltd.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+use crate::util::fxhash::FxBuildHasher;
+use std::collections::HashSet;
+
+pub struct PosMatcher {
+    ids: HashSet<u16, FxBuildHasher>,
+}
+
+impl PosMatcher {
+    pub fn new<I: IntoIterator<Item = u16>>(pos: I) -> Self {
+        let iter = pos.into_iter();
+        let (min_size, max_size) = iter.size_hint();
+        let size = max_size.unwrap_or(min_size);
+        let mut ids = HashSet::with_capacity_and_hasher(size, FxBuildHasher::default());
+        ids.extend(iter);
+        Self { ids }
+    }
+
+    #[inline]
+    pub fn matches_id(&self, pos_id: u16) -> bool {
+        self.ids.contains(&pos_id)
+    }
+
+    pub fn num_entries(&self) -> usize {
+        self.ids.len()
+    }
+
+    pub fn entries(&self) -> impl Iterator<Item = u16> + '_ {
+        self.ids.iter().cloned()
+    }
+}


Test patch:
diff --git a/python/tests/test_pos_matcher.py b/python/tests/test_pos_matcher.py
new file mode 100644
index 00000000..f0e0d897
--- /dev/null
+++ b/python/tests/test_pos_matcher.py
@@ -0,0 +1,66 @@
+#   Copyright (c) 2021 Works Applications Co., Ltd.
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+import os
+import unittest
+
+import sudachipy
+from sudachipy import Dictionary
+
+
+class PosMatcherTestCase(unittest.TestCase):
+
+    def setUp(self):
+        resource_dir = os.path.join(os.path.dirname(
+            os.path.abspath(__file__)), 'resources')
+        self.dict = Dictionary(os.path.join(
+            resource_dir, 'sudachi.json'), resource_dir)
+        self.tokenizer_obj = self.dict.create()
+
+    def test_create_fn(self):
+        m = self.dict.pos_matcher(lambda p: p[0] == "名詞")
+        self.assertIsNotNone(m)
+
+    def test_create_declarative(self):
+        m = self.dict.pos_matcher([("名詞",)])
+        self.assertEqual(list(m), [('名詞', '固有名詞', '地名', '一般', '*', '*'),
+                                   ('名詞', '普通名詞', '一般', '*', '*', '*'),
+                                   ('名詞', '数詞', '*', '*', '*', '*')])
+
+    def test_create_declarative_empty(self):
+        m = self.dict.pos_matcher([()])
+        self.assertEqual(len(m), 9)
+
+    def test_iter(self):
+        m = self.dict.pos_matcher(lambda p: p[0] != "名詞")
+        self.assertEqual(list(m), [('助動詞', '*', '*', '*', '助動詞-タ', '終止形-一般'),
+                                   ('助詞', '接続助詞', '*', '*', '*', '*'),
+                                   ('助詞', '格助詞', '*', '*', '*', '*'),
+                                   ('動詞', '非自立可能', '*', '*', '五段-カ行', '終止形-一般'),
+                                   ('動詞', '非自立可能', '*', '*', '五段-カ行', '連用形-促音便'),
+                                   ('被子植物門', '双子葉植物綱', 'ムクロジ目', 'ミカン科', 'ミカン属', 'スダチ')])
+
+    def test_skips(self):
+        m = self.dict.pos_matcher([(None, None, None, None, None, "終止形-一般")])
+        self.assertEqual(list(m), [('助動詞', '*', '*', '*', '助動詞-タ', '終止形-一般'),
+                                   ('動詞', '非自立可能', '*', '*', '五段-カ行', '終止形-一般')])
+
+    def test_match_nouns(self):
+        tok = self.dict.create(mode=sudachipy.SplitMode.A)
+        m = self.dict.pos_matcher([("名詞",)])
+        data = [x.surface() for x in tok.tokenize("東京に行く") if m(x)]
+        self.assertEqual(data, ["東京"])
+
+
+if __name__ == '__main__':
+    unittest.main()