Merge pull request #1263 from PyThaiNLP/copilot/add-type-hints-to-tokenize

Add 100% type hints to pythainlp.tokenize submodule

Author: bact

pyproject.toml

@@ -411,7 +411,6 @@ module = [
     "khamyo.*",
     "khanaa.*",
     "multiel.*",
-    "nlpo3.*",
     "nltk.*",
     "numpy.*",
     "onnxruntime.*",

pythainlp/tokenize/_utils.py

@@ -7,13 +7,13 @@
 from __future__ import annotations
 
 import re
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 
 _DIGITS_WITH_SEPARATOR = re.compile(r"(\d+[\.\,:])+\d+")
 
 
 def apply_postprocessors(
-    segments: list[str], postprocessors: list[Callable[[list[str]], list[str]]]
+    segments: list[str], postprocessors: Sequence[Callable[[list[str]], list[str]]]
 ) -> list[str]:
     """A list of callables to apply to a raw segmentation result.
     """

pythainlp/tokenize/attacut.py

@@ -10,6 +10,7 @@
 from __future__ import annotations
 
 import threading
+from typing import cast
 
 from attacut import Tokenizer
 
@@ -24,7 +25,7 @@ def __init__(self, model="attacut-sc"):
         self._tokenizer = Tokenizer(model=self._MODEL_NAME)
 
     def tokenize(self, text: str) -> list[str]:
-        return self._tokenizer.tokenize(text)
+        return cast(list[str], self._tokenizer.tokenize(text))
 
 
 _tokenizers: dict[str, AttacutTokenizer] = {}

pythainlp/tokenize/budoux.py

@@ -13,6 +13,7 @@
 from __future__ import annotations
 
 import threading
+from typing import cast
 
 _parser = None
 _parser_lock = threading.Lock()
@@ -55,6 +56,6 @@ def segment(text: str) -> list[str]:
             _parser = _init_parser()
         parser = _parser
 
-    result = parser.parse(text)
+    result = cast(list[str], parser.parse(text))
 
     return result

pythainlp/tokenize/core.py

@@ -32,7 +32,7 @@
 
 def word_detokenize(
     segments: Union[list[list[str]], list[str]], output: str = "str"
-) -> Union[list[str], str]:
+) -> Union[list[list[str]], str]:
     """Word detokenizer.
 
     Detokenizes the list of words in each sentence into text.
@@ -49,18 +49,18 @@ def word_detokenize(
         print(word_detokenize(["เรา", "เล่น"]))
         # output: เราเล่น
     """
-    list_all = []
+    list_all: list[list[str]] = []
 
     if isinstance(segments[0], str):
-        segments = [segments]
+        segments = [segments]  # type: ignore[assignment]
 
     from pythainlp import thai_characters
 
     for i, s in enumerate(segments):
-        list_sents = []
-        add_index = []
-        space_index = []
-        mark_index = []
+        list_sents: list[str] = []
+        add_index: list[int] = []
+        space_index: list[int] = []
+        mark_index: list[int] = []
         for j, w in enumerate(s):
             if j > 0:
                 # previous word
@@ -92,9 +92,9 @@ def word_detokenize(
     if output == "list":
         return list_all
 
-    text = []
-    for i in list_all:
-        text.append("".join(i))
+    text: list[str] = []
+    for sent_tokens in list_all:
+        text.append("".join(sent_tokens))
     return " ".join(text)
 
 
@@ -259,56 +259,56 @@ def word_tokenize(
 
         segments = segment(text, custom_dict, safe_mode=True)
     elif engine == "attacut":
-        from pythainlp.tokenize.attacut import segment
+        from pythainlp.tokenize.attacut import segment as attacut_segment  # noqa: I001
 
-        segments = segment(text)
+        segments = attacut_segment(text)
     elif engine == "longest":
-        from pythainlp.tokenize.longest import segment
+        from pythainlp.tokenize.longest import segment as longest_segment  # noqa: I001
 
-        segments = segment(text, custom_dict)
+        segments = longest_segment(text, custom_dict)
     elif engine in ("mm", "multi_cut"):
-        from pythainlp.tokenize.multi_cut import segment
+        from pythainlp.tokenize.multi_cut import segment as multi_cut_segment  # noqa: I001
 
-        segments = segment(text, custom_dict)
+        segments = multi_cut_segment(text, custom_dict)
     elif engine == "deepcut":  # deepcut can optionally use dictionary
-        from pythainlp.tokenize.deepcut import segment
+        from pythainlp.tokenize.deepcut import segment as deepcut_segment  # noqa: I001
 
         if custom_dict:
-            custom_dict = list(custom_dict)
-            segments = segment(text, custom_dict)
+            custom_dict = list(custom_dict)  # type: ignore[assignment]
+            segments = deepcut_segment(text, custom_dict)
         else:
-            segments = segment(text)
+            segments = deepcut_segment(text)
     elif engine == "icu":
-        from pythainlp.tokenize.pyicu import segment
+        from pythainlp.tokenize.pyicu import segment as pyicu_segment  # noqa: I001
 
-        segments = segment(text)
+        segments = pyicu_segment(text)
     elif engine == "budoux":
-        from pythainlp.tokenize.budoux import segment
+        from pythainlp.tokenize.budoux import segment as budoux_segment  # noqa: I001
 
-        segments = segment(text)
+        segments = budoux_segment(text)
     elif engine == "nercut":
-        from pythainlp.tokenize.nercut import segment
+        from pythainlp.tokenize.nercut import segment as nercut_segment  # noqa: I001
 
-        segments = segment(text)
+        segments = nercut_segment(text)
     elif engine == "sefr_cut":
-        from pythainlp.tokenize.sefr_cut import segment
+        from pythainlp.tokenize.sefr_cut import segment as sefrcut_segment  # noqa: I001
 
-        segments = segment(text)
+        segments = sefrcut_segment(text)
     elif engine == "tltk":
-        from pythainlp.tokenize.tltk import segment
+        from pythainlp.tokenize.tltk import segment as tltk_segment  # noqa: I001
 
-        segments = segment(text)
+        segments = tltk_segment(text)
     elif engine == "oskut":
-        from pythainlp.tokenize.oskut import segment
+        from pythainlp.tokenize.oskut import segment as oskut_segment  # noqa: I001
 
-        segments = segment(text)
+        segments = oskut_segment(text)
     elif engine == "nlpo3":
-        from pythainlp.tokenize.nlpo3 import segment
+        from pythainlp.tokenize.nlpo3 import segment as nlpo3_segment  # noqa: I001
 
         # Currently cannot handle custom_dict from inside word_tokenize(),
         # due to difference in type.
         # if isinstance(custom_dict, str):
-        #    segments = segment(text, custom_dict=custom_dict)
+        #    segments = nlpo3_segment(text, custom_dict=custom_dict)
         # elif not isinstance(custom_dict, str) and not custom_dict:
         #    raise ValueError(
         #        f"""Tokenizer \"{engine}\":
@@ -317,8 +317,8 @@ def word_tokenize(
         #        See pythainlp.tokenize.nlpo3.load_dict()"""
         #    )
         # else:
-        #    segments = segment(text)
-        segments = segment(text)
+        #    segments = nlpo3_segment(text)
+        segments = nlpo3_segment(text)
     else:
         raise ValueError(
             f"""Tokenizer \"{engine}\" not found.
@@ -413,7 +413,7 @@ def sent_tokenize(
     text: Union[str, list[str]],
     engine: str = DEFAULT_SENT_TOKENIZE_ENGINE,
     keep_whitespace: bool = True,
-) -> list[str]:
+) -> Union[list[str], list[list[str]]]:
     """Sentence tokenizer.
 
     Tokenizes running text into "sentences". Supports both string and list of strings.
@@ -632,7 +632,7 @@ def paragraph_tokenize(
             It might be a typo; if not, please consult our document."""
         )
 
-    return segments
+    return segments  # type: ignore[return-value]
 
 
 def subword_tokenize(
@@ -719,36 +719,41 @@ def subword_tokenize(
     segments = []
 
     if engine == "tcc":
-        from pythainlp.tokenize.tcc import segment
+        from pythainlp.tokenize.tcc import segment as tcc_segment
+        segments = tcc_segment(text)
     elif engine == "tcc_p":
-        from pythainlp.tokenize.tcc_p import segment
+        from pythainlp.tokenize.tcc_p import segment as tcc_p_segment
+        segments = tcc_p_segment(text)
     elif engine == "etcc":
-        from pythainlp.tokenize.etcc import segment
+        from pythainlp.tokenize.etcc import segment as etcc_segment
+        segments = etcc_segment(text)
     elif engine == "wangchanberta":
-        from pythainlp.wangchanberta import segment
+        from pythainlp.wangchanberta import segment as wangchanberta_segment
+        segments = wangchanberta_segment(text)
     elif engine == "dict":  # use syllable dictionary
         words = word_tokenize(text)
         for word in words:
             segments.extend(
                 word_tokenize(text=word, custom_dict=syllable_dict_trie())
             )
     elif engine == "ssg":
-        from pythainlp.tokenize.ssg import segment
+        from pythainlp.tokenize.ssg import segment as ssg_segment
+        segments = ssg_segment(text)
     elif engine == "tltk":
-        from pythainlp.tokenize.tltk import syllable_tokenize as segment
+        from pythainlp.tokenize.tltk import syllable_tokenize as tltk_segment
+        segments = tltk_segment(text)
     elif engine == "han_solo":
-        from pythainlp.tokenize.han_solo import segment
+        from pythainlp.tokenize.han_solo import segment as han_solo_segment
+        segments = han_solo_segment(text)
     elif engine == "phayathai":
-        from pythainlp.phayathaibert import segment
+        from pythainlp.phayathaibert import segment as phayathai_segment
+        segments = phayathai_segment(text)
     else:
         raise ValueError(
             f"""Tokenizer \"{engine}\" not found.
             It might be a typo; if not, please consult our document."""
         )
 
-    if not segments:
-        segments = segment(text)
-
     if not keep_whitespace:
         segments = strip_whitespace(segments)
 

pythainlp/tokenize/crfcut.py

@@ -183,10 +183,7 @@ def segment(text: str) -> list[str]:
     :param str text: text to be tokenized into sentences
     :return: list of words, tokenized from the text
     """
-    if isinstance(text, str):
-        toks = word_tokenize(text)
-    else:
-        toks = text
+    toks = word_tokenize(text)
     feat = extract_features(toks)
     labs = _tagger.tag(feat)
     labs[-1] = "E"  # make sure it cuts the last sentence

pythainlp/tokenize/deepcut.py

@@ -12,7 +12,7 @@
 
 from __future__ import annotations
 
-from typing import Union
+from typing import Union, cast
 
 try:
     from deepcut import tokenize
@@ -29,6 +29,6 @@ def segment(text: str, custom_dict: Union[Trie, list[str], str] = []) -> list[st
         if isinstance(custom_dict, Trie):
             custom_dict = list(custom_dict)
 
-        return tokenize(text, custom_dict)
+        return cast(list[str], tokenize(text, custom_dict))
 
-    return tokenize(text)
+    return cast(list[str], tokenize(text))

pythainlp/tokenize/longest.py

@@ -106,7 +106,7 @@ def __longest_matching(self, text: str, begin_pos: int) -> str:
         else:
             return ""
 
-    def __segment(self, text: str):
+    def __segment(self, text: str) -> list[str]:
         begin_pos = 0
         len_text = len(text)
         tokens: list[str] = []

pythainlp/tokenize/nercut.py

@@ -40,7 +40,7 @@ def segment(
     :param class tagger: NER tagger engine
     :return: list of words, tokenized from the text
     """
-    if not isinstance(text, str):
+    if not text:
         return []
 
     tagged_words = tagger.tag(text, pos=False)
@@ -73,7 +73,5 @@ def segment(
                 words.append(combining_word)
             elif curr_tag.startswith("I-") and combining_word != "":
                 words.append(combining_word)
-            else:
-                pass
 
     return words

pythainlp/tokenize/nlpo3.py

@@ -6,9 +6,11 @@
 import threading
 from importlib.resources import as_file, files
 from sys import stderr
+from typing import TYPE_CHECKING
 
-from nlpo3 import load_dict as nlpo3_load_dict
-from nlpo3 import segment as nlpo3_segment
+if TYPE_CHECKING:
+    from nlpo3 import load_dict as nlpo3_load_dict  # noqa: F401
+    from nlpo3 import segment as nlpo3_segment  # noqa: F401
 
 from pythainlp.corpus.common import _THAI_WORDS_FILENAME
 
@@ -25,6 +27,13 @@ def _ensure_default_dict_loaded():
     The context manager is kept alive for the lifetime of the program
     to prevent cleanup of temporary files while the dictionary is in use.
     """
+    try:
+        from nlpo3 import load_dict as nlpo3_load_dict
+    except ImportError as ex:
+        raise ImportError(
+            "nlpo3 is not installed. Install it with: pip install nlpo3"
+        ) from ex
+
     global _NLPO3_DEFAULT_DICT, _dict_file_ctx
     if _NLPO3_DEFAULT_DICT is None:
         with _load_lock:
@@ -57,6 +66,13 @@ def load_dict(file_path: str, dict_name: str) -> bool:
         * \
             https://github.com/PyThaiNLP/nlpo3
     """
+    try:
+        from nlpo3 import load_dict as nlpo3_load_dict
+    except ImportError as ex:
+        raise ImportError(
+            "nlpo3 is not installed. Install it with: pip install nlpo3"
+        ) from ex
+
     msg, success = nlpo3_load_dict(file_path=file_path, dict_name=dict_name)
     if not success:
         print(msg, file=stderr)
@@ -87,6 +103,13 @@ def segment(
         * \
             https://github.com/PyThaiNLP/nlpo3
     """
+    try:
+        from nlpo3 import segment as nlpo3_segment
+    except ImportError as ex:
+        raise ImportError(
+            "nlpo3 is not installed. Install it with: pip install nlpo3"
+        ) from ex
+
     # Ensure default dict is loaded if it's being used
     if custom_dict == _NLPO3_DEFAULT_DICT_NAME:
         _ensure_default_dict_loaded()