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

Issue:
PosMatcher ergonomics
add set operations to PosMatcher:

* `union` which returns a union of two PosMatchers (operation A and B)
* `intersection` which returns an intersection of two PosMatchers (operation A or B)
* `negation` which returns a negation of a PosMatcher (operation `not`)


Discussion / hints:

Fix patch:
diff --git a/python/py_src/sudachipy/sudachipy.pyi b/python/py_src/sudachipy/sudachipy.pyi
index e7e79089..484d42b0 100644
--- a/python/py_src/sudachipy/sudachipy.pyi
+++ b/python/py_src/sudachipy/sudachipy.pyi
@@ -1,8 +1,5 @@
 from typing import ClassVar, Iterator, List, Tuple, Union, Callable, Iterable, Optional, Literal, Set
 
-import sudachipy
-from sudachipy.pretokenizer import SudachiPreTokenizer
-
 POS = Tuple[str, str, str, str, str, str]
 # POS element
 PE = Optional[str]
@@ -19,6 +16,23 @@ PartialPOS = Union[
 FieldSet = Optional[Set[Literal["surface", "pos", "normalized_form", "dictionary_form", "reading_form",
                                 "word_structure", "split_a", "split_b", "synonym_group_id"]]]
 
+class SplitMode:
+    """
+    Unit to split text.
+
+    A == short mode
+
+    B == middle mode
+
+    C == long mode
+    """
+
+    A: ClassVar[SplitMode] = ...
+    B: ClassVar[SplitMode] = ...
+    C: ClassVar[SplitMode] = ...
+    @classmethod
+    def __init__(cls) -> None: ...
+
 
 class Dictionary:
     """
@@ -48,7 +62,7 @@ class Dictionary:
         ...
 
     def create(self,
-               mode: SplitMode = sudachipy.SplitMode.C,
+               mode: SplitMode = SplitMode.C,
                fields: FieldSet = None) -> Tokenizer:
         """
         Creates a Sudachi Tokenizer.
@@ -61,7 +75,7 @@ class Dictionary:
 
     def pos_matcher(self, target: Union[Iterable[PartialPOS], Callable[[POS], bool]]) -> PosMatcher:
         """
-        Creates a new POS matcher
+        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.
@@ -76,9 +90,9 @@ class Dictionary:
         ...
 
     def pre_tokenizer(self,
-                      mode: SplitMode = sudachipy.SplitMode.C,
+                      mode: SplitMode = SplitMode.C,
                       fields: FieldSet = None,
-                      handler: Optional[Callable[[int, object, MorphemeList], list]] = None) -> SudachiPreTokenizer:
+                      handler: Optional[Callable[[int, object, MorphemeList], list]] = None) -> object:
         """
         Creates HuggingFace Tokenizers-compatible PreTokenizer.
         Requires package `tokenizers` to be installed.
@@ -91,12 +105,20 @@ class Dictionary:
         """
         ...
 
+    def pos_of(self, pos_id: int) -> Optional[POS]:
+        """
+        Returns POS with the given id.
+
+        :param pos_id: POS id
+        :return: POS tuple with the given id.
+        """
+        ...
+
 
 class Morpheme:
     """
     A morpheme (basic semantic unit of language).
     """
-    @classmethod
     def __init__(self) -> None: ...
 
     def begin(self) -> int:
@@ -200,12 +222,13 @@ class Morpheme:
 class MorphemeList:
     """
     A list of morphemes.
+    An object can not be instantiated manually.
+    Use Tokenizer.tokenize("") to create an empty morpheme list.
     """
-    @classmethod
     def __init__(self) -> None: ...
 
     @classmethod
-    def empty(dict) -> MorphemeList:
+    def empty(cls, dict) -> MorphemeList:
         """
         Returns an empty morpheme list with dictionary.
         """
@@ -228,31 +251,16 @@ class MorphemeList:
     def __len__(self) -> int: ...
 
 
-class SplitMode:
-    """
-    Unit to split text.
-
-    A == short mode
-
-    B == middle mode
-
-    C == long mode
-    """
 
-    A: ClassVar[SplitMode] = ...
-    B: ClassVar[SplitMode] = ...
-    C: ClassVar[SplitMode] = ...
-    @classmethod
-    def __init__(cls) -> None: ...
 
 
 class Tokenizer:
-    SplitMode: ClassVar[sudachipy.SplitMode] = ...
+    SplitMode: ClassVar[SplitMode] = ...
     @classmethod
     def __init__(cls) -> None: ...
 
     def tokenize(self, text: str,
-                 mode: sudachipy.SplitMode = ...,
+                 mode: SplitMode = ...,
                  out: Optional[MorphemeList] = None) -> MorphemeList:
         """
         Break text into morphemes.
@@ -295,4 +303,32 @@ class PosMatcher:
         :param m: morpheme
         :return: if morpheme has matching POS
         """
+        ...
+
+    def __or__(self, other: PosMatcher) -> PosMatcher:
+        """
+        Returns a POS matcher which matches a POS if any of two matchers would match it
+        :return: PosMatcher
+        """
+        ...
+
+    def __and__(self, other: PosMatcher) -> PosMatcher:
+        """
+        Returns a POS matcher which matches a POS if both matchers would match it at the same time
+        :return: PosMatcher
+        """
+        ...
+
+    def __sub__(self, other: PosMatcher) -> PosMatcher:
+        """
+        Returns a POS matcher which matches a POS if self would match the POS and other would not match the POS
+        :return: PosMatcher
+        """
+        ...
+
+    def __invert__(self) -> PosMatcher:
+        """
+        Returns a POS matcher which matches all POS tags except ones defined in the current POS matcher
+        :return: PosMatcher
+        """
         ...
\ No newline at end of file
diff --git a/python/src/dictionary.rs b/python/src/dictionary.rs
index 09dd4e72..35206887 100644
--- a/python/src/dictionary.rs
+++ b/python/src/dictionary.rs
@@ -277,6 +277,13 @@ impl PyDictionary {
         self.dictionary = None;
     }
 
+    /// Get POS Tuple by its id
+    #[pyo3(text_signature = "($self, pos_id: int)")]
+    fn pos_of<'p>(&'p self, py: Python<'p>, pos_id: usize) -> Option<&'p PyTuple> {
+        let dic = self.dictionary.as_ref().unwrap();
+        dic.pos.get(pos_id).map(|x| x.as_ref(py))
+    }
+
     fn __repr__(&self) -> PyResult<String> {
         config_repr(&self.config).map_err(|e| {
             return PyException::new_err(format!("{:?}", e));
diff --git a/python/src/pos_matcher.rs b/python/src/pos_matcher.rs
index 8ac83d06..359d8b4b 100644
--- a/python/src/pos_matcher.rs
+++ b/python/src/pos_matcher.rs
@@ -14,15 +14,18 @@
  *  limitations under the License.
  */
 
-use crate::dictionary::PyDicData;
-use crate::morpheme::PyMorpheme;
+use std::sync::Arc;
+
 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;
 
+use crate::dictionary::PyDicData;
+use crate::morpheme::PyMorpheme;
+
 #[pyclass(name = "PosMatcher", module = "sudachipy")]
 pub struct PyPosMatcher {
     matcher: PosMatcher,
@@ -131,6 +134,59 @@ impl PyPosMatcher {
     pub fn __len__(&self) -> usize {
         self.matcher.num_entries()
     }
+
+    pub fn __or__(&self, other: &Self) -> Self {
+        assert_eq!(
+            Arc::as_ptr(&self.dic),
+            Arc::as_ptr(&other.dic),
+            "incompatible dictionaries"
+        );
+        let matcher = self.matcher.union(&other.matcher);
+        Self {
+            dic: self.dic.clone(),
+            matcher,
+        }
+    }
+
+    pub fn __and__(&self, other: &Self) -> Self {
+        assert_eq!(
+            Arc::as_ptr(&self.dic),
+            Arc::as_ptr(&other.dic),
+            "incompatible dictionaries"
+        );
+        let matcher = self.matcher.intersection(&other.matcher);
+        Self {
+            dic: self.dic.clone(),
+            matcher,
+        }
+    }
+
+    pub fn __sub__(&self, other: &Self) -> Self {
+        assert_eq!(
+            Arc::as_ptr(&self.dic),
+            Arc::as_ptr(&other.dic),
+            "incompatible dictionaries"
+        );
+        let matcher = self.matcher.difference(&other.matcher);
+        Self {
+            dic: self.dic.clone(),
+            matcher,
+        }
+    }
+
+    pub fn __invert__(&self) -> Self {
+        let max_id = self.dic.pos.len();
+        // map -> filter chain is needed to handle exactly u16::MAX POS entries
+        let values = (0..max_id)
+            .into_iter()
+            .map(|x| x as u16)
+            .filter(|id| !self.matcher.matches_id(*id));
+        let matcher = PosMatcher::new(values);
+        Self {
+            matcher,
+            dic: self.dic.clone(),
+        }
+    }
 }
 
 #[pyclass(name = "PosMatcherIterator", module = "sudachipy")]
diff --git a/sudachi/src/pos.rs b/sudachi/src/pos.rs
index 2136a485..7862a75b 100644
--- a/sudachi/src/pos.rs
+++ b/sudachi/src/pos.rs
@@ -43,4 +43,22 @@ impl PosMatcher {
     pub fn entries(&self) -> impl Iterator<Item = u16> + '_ {
         self.ids.iter().cloned()
     }
+
+    pub fn union(&self, other: &PosMatcher) -> PosMatcher {
+        let mut ids = self.ids.clone();
+        ids.extend(other.ids.iter());
+        PosMatcher { ids }
+    }
+
+    pub fn intersection(&self, other: &PosMatcher) -> PosMatcher {
+        let mut ids = self.ids.clone();
+        ids.retain(|id| other.ids.contains(id));
+        PosMatcher { ids }
+    }
+
+    pub fn difference(&self, other: &PosMatcher) -> PosMatcher {
+        let mut ids = self.ids.clone();
+        ids.retain(|id| !other.ids.contains(id));
+        PosMatcher { ids }
+    }
 }


Test patch:
diff --git a/python/tests/test_dictionary.py b/python/tests/test_dictionary.py
index a1150436..94617cd1 100644
--- a/python/tests/test_dictionary.py
+++ b/python/tests/test_dictionary.py
@@ -32,6 +32,9 @@ def tearDown(self) -> None:
     def test_create(self):
         self.assertEqual(Tokenizer, type(self.dict_.create()))
 
+    def test_pos_of(self):
+        self.assertIsNotNone(self.dict_.pos_of(0))
+
     def test_repr(self):
         repr_str = repr(self.dict_)
         self.assertTrue(repr_str.startswith("<SudachiDictionary(system="))
diff --git a/python/tests/test_pos_matcher.py b/python/tests/test_pos_matcher.py
index f0e0d897..c5ecd6a4 100644
--- a/python/tests/test_pos_matcher.py
+++ b/python/tests/test_pos_matcher.py
@@ -61,6 +61,31 @@ def test_match_nouns(self):
         data = [x.surface() for x in tok.tokenize("東京に行く") if m(x)]
         self.assertEqual(data, ["東京"])
 
+    def test_union(self):
+        m1 = self.dict.pos_matcher(lambda p: p[0] == "名詞")
+        m2 = self.dict.pos_matcher(lambda p: p[0] == "動詞")
+        m3 = m1 | m2
+        self.assertEqual(len(m1) + len(m2), len(m3))
+
+    def test_intersection(self):
+        m1 = self.dict.pos_matcher(lambda p: p[5] == "終止形-一般")
+        m2 = self.dict.pos_matcher(lambda p: p[0] == "動詞")
+        m3 = m1 & m2
+        self.assertEqual(1, len(m3))
+        self.assertEqual(list(m3), [('動詞', '非自立可能', '*', '*', '五段-カ行', '終止形-一般')])
+
+    def test_difference(self):
+        m1 = self.dict.pos_matcher(lambda p: p[5] == "終止形-一般")
+        m2 = self.dict.pos_matcher(lambda p: p[0] == "動詞")
+        m3 = m1 - m2
+        self.assertEqual(1, len(m3))
+        self.assertEqual(list(m3), [('助動詞', '*', '*', '*', '助動詞-タ', '終止形-一般')])
+
+    def test_invert(self):
+        m1 = self.dict.pos_matcher(lambda p: p[5] == "終止形-一般")
+        self.assertEqual(len(m1), 2)
+        self.assertEqual(len(~m1), 7)
+
 
 if __name__ == '__main__':
     unittest.main()