Merge pull request #926 from WTFPUn/tone_collate

Tone collate fixed

Author: wannaphong

pythainlp/util/collate.py

@@ -15,7 +15,9 @@
 def _thkey(word: str) -> str:
     cv = _RE_TONE.sub("", word)  # remove tone
     cv = _RE_LV_C.sub("\\2\\1", cv)  # switch lead vowel
-    tone = _RE_TONE.sub(" ", word)  # just tone
+
+    tone_match = _RE_TONE.search(word)
+    tone = tone_match.group() if tone_match else ""
     return cv + tone
 
 

tests/test_util.py

@@ -79,6 +79,14 @@ def test_collate(self):
             collate(["ไก่", "เป็ด", "หมู", "วัว"]),
             ["ไก่", "เป็ด", "วัว", "หมู"],
         )
+        self.assertEqual(
+            collate(["ก้วย", "ก๋วย", "กวย", "ก่วย", "ก๊วย"]),
+            collate(["ก๋วย", "ก่วย", "ก้วย", "ก๊วย", "กวย"]),
+        )  # should guarantee same order
+        self.assertEqual(
+            collate(["ก้วย", "ก๋วย", "ก่วย", "กวย", "ก้วย", "ก่วย", "ก๊วย"]),
+            ["กวย", "ก่วย", "ก่วย", "ก้วย", "ก้วย", "ก๊วย", "ก๋วย"],
+        )
 
     # ### pythainlp.util.numtoword