Complete type annotations across codebase (97% coverage)

pythainlp/ancient/aksonhan.py

@@ -11,9 +11,11 @@
 from pythainlp.util import Trie
 
 _dict_aksonhan: dict[str, str] = {}
+i: str
 for i in list(thai_consonants):
     if i == "ร":
         continue
+    j: str
     for j in list(thai_tonemarks):
         _dict_aksonhan[i + j + i] = "ั" + j + i
         _dict_aksonhan[i + i + j + i] = i + "ั" + j + i

pythainlp/augment/lm/fasttext.py

@@ -30,13 +30,11 @@ def __init__(self, model_path: str) -> None:
         from gensim.models.keyedvectors import KeyedVectors
 
         if model_path.endswith(".bin"):
-            self.model: Union[FastText, KeyedVectors] = (
-                FastText_gensim.load_facebook_vectors(model_path)
-            )
+            self.model: Union["FastText", "KeyedVectors"] = FastText_gensim.load_facebook_vectors(model_path)
         elif model_path.endswith(".vec"):
-            self.model = KeyedVectors.load_word2vec_format(model_path)
+            self.model: Union["FastText", "KeyedVectors"] = KeyedVectors.load_word2vec_format(model_path)
         else:
-            self.model = FastText_gensim.load(model_path)
+            self.model: Union["FastText", "KeyedVectors"] = FastText_gensim.load(model_path)
         self.dict_wv: list[str] = list(self.model.key_to_index.keys())
 
     def tokenize(self, text: str) -> list[str]:

pythainlp/augment/lm/phayathaibert.py

@@ -28,16 +28,16 @@ def __init__(self) -> None:
             pipeline,
         )
 
-        self.tokenizer = AutoTokenizer.from_pretrained(_MODEL_NAME)  # type: ignore[assignment]
-        self.model_for_masked_lm = AutoModelForMaskedLM.from_pretrained(  # type: ignore[assignment]
+        self.tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(_MODEL_NAME)
+        self.model_for_masked_lm: AutoModelForMaskedLM = AutoModelForMaskedLM.from_pretrained(
             _MODEL_NAME
         )
-        self.model = pipeline(
+        self.model: Pipeline = pipeline(
             "fill-mask",
             tokenizer=self.tokenizer,
             model=self.model_for_masked_lm,
         )
-        self.processor = ThaiTextProcessor()
+        self.processor: ThaiTextProcessor = ThaiTextProcessor()
 
     def generate(
         self,

pythainlp/augment/lm/wangchanberta.py

@@ -25,29 +25,29 @@ def __init__(self) -> None:
             pipeline,
         )
 
-        self.model_name = "airesearch/wangchanberta-base-att-spm-uncased"
-        self.target_tokenizer = CamembertTokenizer
-        self.tokenizer = CamembertTokenizer.from_pretrained(  # type: ignore[assignment]
+        self.model_name: str = "airesearch/wangchanberta-base-att-spm-uncased"
+        self.target_tokenizer: type[CamembertTokenizer] = CamembertTokenizer
+        self.tokenizer: CamembertTokenizer = CamembertTokenizer.from_pretrained(
             self.model_name, revision="main"
         )
         self.tokenizer.additional_special_tokens = [
             "<s>NOTUSED",
             "</s>NOTUSED",
             "<_>",
         ]
-        self.fill_mask = pipeline(
+        self.fill_mask: Pipeline = pipeline(
             task="fill-mask",
             tokenizer=self.tokenizer,
             model=f"{self.model_name}",
             revision="main",
         )
-        self.MASK_TOKEN = self.tokenizer.mask_token
+        self.MASK_TOKEN: str = self.tokenizer.mask_token
 
     def generate(
         self, sentence: str, num_replace_tokens: int = 3
     ) -> list[str]:
         sent2: list[str] = []
-        self.input_text = sentence
+        self.input_text: str = sentence
         sent = [
             i for i in self.tokenizer.tokenize(self.input_text) if i != "▁"
         ]

pythainlp/augment/word2vec/bpemb_wv.py

@@ -32,8 +32,8 @@ def __init__(
     ) -> None:
         from bpemb import BPEmb
 
-        self.bpemb_temp = BPEmb(lang=lang, dim=dim, vs=vs)
-        self.model = self.bpemb_temp.emb
+        self.bpemb_temp: BPEmb = BPEmb(lang=lang, dim=dim, vs=vs)
+        self.model: KeyedVectors = self.bpemb_temp.emb
         self.load_w2v()
 
     def tokenizer(self, text: str) -> list[str]:
@@ -44,7 +44,7 @@ def tokenizer(self, text: str) -> list[str]:
 
     def load_w2v(self) -> None:
         """Load BPEmb model"""
-        self.aug = Word2VecAug(
+        self.aug: Word2VecAug = Word2VecAug(
             self.model, tokenize=self.tokenizer, type="model"
         )
 
@@ -68,11 +68,11 @@ def augment(
             aug.augment("ผมเรียน", n_sent=2, p=0.5)
             # output: ['ผมสอน', 'ผมเข้าเรียน']
         """
-        self.sentence = sentence.replace(" ", "▁")
-        self.temp = self.aug.augment(self.sentence, n_sent, p=p)
-        self.temp_new = []
+        self.sentence: str = sentence.replace(" ", "▁")
+        self.temp: list[tuple[str, ...]] = self.aug.augment(self.sentence, n_sent, p=p)
+        self.temp_new: list[str] = []
         for i in self.temp:
-            self.t = ""
+            self.t: str = ""
             for j in i:
                 self.t += j.replace("▁", "")
             self.temp_new.append(self.t)

pythainlp/augment/word2vec/core.py

@@ -29,15 +29,13 @@ def __init__(
 
         self.tokenizer: Callable[[str], list[str]] = tokenize
         if type == "file":
-            self.model: "KeyedVectors" = (
-                word2vec.KeyedVectors.load_word2vec_format(model)
-            )
+            self.model: "KeyedVectors" = word2vec.KeyedVectors.load_word2vec_format(model)
         elif type == "binary":
-            self.model = word2vec.KeyedVectors.load_word2vec_format(
+            self.model: "KeyedVectors" = word2vec.KeyedVectors.load_word2vec_format(
                 model, binary=True, unicode_errors="ignore"
             )
         else:
-            self.model = model
+            self.model: "KeyedVectors" = model  # type: ignore[assignment]
         self.dict_wv: list[str] = list(self.model.key_to_index.keys())
 
     def modify_sent(self, sent: list[str], p: float = 0.7) -> list[list[str]]:

pythainlp/augment/word2vec/ltw2v.py

@@ -3,15 +3,9 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
+from typing import Optional
 
-if TYPE_CHECKING:
-    from pythainlp.augment.word2vec.core import Word2VecAug
-
-from pythainlp.augment.word2vec.core import Word2VecAug as _Word2VecAug
-
-# Make it accessible for runtime
-Word2VecAug: type[_Word2VecAug] = _Word2VecAug
+from pythainlp.augment.word2vec.core import Word2VecAug
 from pythainlp.corpus import get_corpus_path
 from pythainlp.tokenize import word_tokenize
 
@@ -27,7 +21,7 @@ class LTW2VAug:
     aug: Word2VecAug
 
     def __init__(self) -> None:
-        self.ltw2v_wv = get_corpus_path("ltw2v")
+        self.ltw2v_wv: Optional[str] = get_corpus_path("ltw2v")
         self.load_w2v()
 
     def tokenizer(self, text: str) -> list[str]:
@@ -43,7 +37,7 @@ def load_w2v(self) -> None:  # insert substitute
                 "LTW2V word2vec model not found. "
                 "Please download it first using pythainlp.corpus.download('ltw2v_wv')"
             )
-        self.aug = Word2VecAug(self.ltw2v_wv, self.tokenizer, type="binary")
+        self.aug: Word2VecAug = Word2VecAug(self.ltw2v_wv, self.tokenizer, type="binary")
 
     def augment(
         self, sentence: str, n_sent: int = 1, p: float = 0.7

pythainlp/augment/word2vec/thai2fit.py

@@ -3,15 +3,9 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
+from typing import Optional
 
-if TYPE_CHECKING:
-    from pythainlp.augment.word2vec.core import Word2VecAug
-
-from pythainlp.augment.word2vec.core import Word2VecAug as _Word2VecAug
-
-# Make it accessible for runtime
-Word2VecAug: type[_Word2VecAug] = _Word2VecAug
+from pythainlp.augment.word2vec.core import Word2VecAug
 from pythainlp.corpus import get_corpus_path
 from pythainlp.tokenize import thai2fit_tokenizer
 
@@ -27,15 +21,15 @@ class Thai2fitAug:
     aug: Word2VecAug
 
     def __init__(self) -> None:
-        self.thai2fit_wv = get_corpus_path("thai2fit_wv")
+        self.thai2fit_wv: Optional[str] = get_corpus_path("thai2fit_wv")
         self.load_w2v()
 
     def tokenizer(self, text: str) -> list[str]:
         """:param str text: Thai text
         :rtype: List[str]
         """
         tok = thai2fit_tokenizer()
-        return tok.word_tokenize(text)  # type: ignore[no-any-return]
+        return tok.word_tokenize(text)
 
     def load_w2v(self) -> None:
         """Load Thai2Fit's word2vec model"""
@@ -44,7 +38,7 @@ def load_w2v(self) -> None:
                 "Thai2Fit word2vec model not found. "
                 "Please download it first using pythainlp.corpus.download('thai2fit_wv')"
             )
-        self.aug = Word2VecAug(self.thai2fit_wv, self.tokenizer, type="binary")
+        self.aug: Word2VecAug = Word2VecAug(self.thai2fit_wv, self.tokenizer, type="binary")
 
     def augment(
         self, sentence: str, n_sent: int = 1, p: float = 0.7

pythainlp/augment/wordnet.py

@@ -12,7 +12,7 @@
 
 import itertools
 from collections import OrderedDict
-from typing import Callable, Optional
+from typing import Any, Callable, Optional
 
 from nltk.corpus import wordnet as wn
 
@@ -112,12 +112,24 @@ def postype2wordnet(pos: str, corpus: str) -> Optional[str]:
     """
     if corpus not in ["orchid"]:
         return None
-    return orchid[pos]  # type: ignore[no-any-return]
+    return orchid[pos]
 
 
 class WordNetAug:
     """Text Augment using wordnet"""
 
+    synonyms: list[str]
+    list_synsets: list
+    p2w_pos: Optional[str]
+    synset: Any
+    syn: str
+    synonyms_without_duplicates: list[str]
+    list_words: list[str]
+    list_synonym: list
+    p_all: int
+    list_pos: list[tuple[str, str]]
+    temp: list[str]
+
     def __init__(self) -> None:
         pass
 
@@ -137,13 +149,13 @@ def find_synonyms(
         """
         self.synonyms: list[str] = []
         if pos is None:
-            self.list_synsets = wordnet.synsets(word)
+            self.list_synsets: list = wordnet.synsets(word)
         else:
-            self.p2w_pos = postype2wordnet(pos, postag_corpus)
+            self.p2w_pos: Optional[str] = postype2wordnet(pos, postag_corpus)
             if self.p2w_pos != "":
-                self.list_synsets = wordnet.synsets(word, pos=self.p2w_pos)
+                self.list_synsets: list = wordnet.synsets(word, pos=self.p2w_pos)
             else:
-                self.list_synsets = wordnet.synsets(word)
+                self.list_synsets: list = wordnet.synsets(word)
 
         for self.synset in wordnet.synsets(word):
             for self.syn in self.synset.lemma_names(lang="tha"):
@@ -189,7 +201,7 @@ def augment(
         """
         new_sentences = []
         self.list_words: list[str] = tokenize(sentence)
-        self.list_synonym: list[list[str]] = []
+        self.list_synonym: list = []
         self.p_all: int = 1
         if postag:
             self.list_pos: list[tuple[str, str]] = pos_tag(
@@ -206,7 +218,7 @@ def augment(
                     self.p_all *= len(self.temp)
         else:
             for word in self.list_words:
-                self.temp = self.find_synonyms(word)
+                self.temp: list[str] = self.find_synonyms(word)
                 if not self.temp:
                     self.list_synonym.append([word])
                 else:

pythainlp/chat/core.py

@@ -3,13 +3,16 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 if TYPE_CHECKING:
     import torch
 
 
 class ChatBotModel:
+    history: list[tuple[str, str]]
+    model: Any
+
     def __init__(self) -> None:
         """Chat using AI generation"""
         self.history: list[tuple[str, str]] = []
@@ -46,7 +49,7 @@ def load_model(
         if model_name == "wangchanglm":
             from pythainlp.generate.wangchanglm import WangChanGLM
 
-            self.model = WangChanGLM()
+            self.model: Any = WangChanGLM()
             self.model.load_model(
                 model_path="pythainlp/wangchanglm-7.5B-sft-en-sharded",
                 return_dict=return_dict,
@@ -94,4 +97,4 @@ def chat(self, text: str) -> str:
         )
         _bot = self.model.gen_instruct(_temp)
         self.history.append((text, _bot))
-        return _bot  # type: ignore[no-any-return]
+        return _bot

pythainlp/classify/param_free.py

@@ -35,8 +35,8 @@ def __init__(
         if model_path:
             self.load(model_path)
         else:
-            self.training_data = np.array(training_data)
-            self.cx2_list = self.train()
+            self.training_data: "NDArray[Any]" = np.array(training_data)
+            self.cx2_list: list[int] = self.train()
 
     def train(self) -> list[int]:
         temp_list = []
@@ -112,5 +112,5 @@ def load(self, path: str) -> None:
 
         with open(path, "r", encoding="utf-8") as f:
             data = json.load(f)
-            self.cx2_list = data["cx2_list"]
-            self.training_data = np.array(data["training_data"])
+            self.cx2_list: list[int] = data["cx2_list"]
+            self.training_data: "NDArray[Any]" = np.array(data["training_data"])

pythainlp/cli/__init__.py

@@ -15,8 +15,8 @@
 if TYPE_CHECKING:
     from types import ModuleType
 
-sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")  # type: ignore[assignment]
-sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8")  # type: ignore[assignment]
+sys.stdout: io.TextIOWrapper = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")
+sys.stderr: io.TextIOWrapper = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8")
 
 # a command should start with a verb when possible
 COMMANDS: list[str] = sorted(

pythainlp/cli/tag.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import argparse
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from pythainlp import cli
 from pythainlp.tag import pos_tag
@@ -38,7 +38,7 @@ def __init__(self, name: str, argv: Sequence[str]) -> None:
         )
 
         args = parser.parse_args(argv)
-        self.args = args
+        self.args: Any = args
 
         tokens = args.text.split(args.separator)
         result = self.run(tokens)
@@ -52,8 +52,8 @@ class POSTaggingApp(SubAppBase):
     run: Callable[[list[str]], list[tuple[str, str]]]
 
     def __init__(self, *args: str, **kwargs: str) -> None:
-        self.separator = "|"
-        self.run = pos_tag
+        self.separator: str = "|"
+        self.run: Callable[[list[str]], list[tuple[str, str]]] = pos_tag
 
         super().__init__(*args, **kwargs)
 

pythainlp/cli/tokenize.py

@@ -74,7 +74,7 @@ def __init__(self, name: str, argv: Sequence[str]) -> None:
         parser.set_defaults(keep_whitespace=True)
 
         args = parser.parse_args(argv)
-        self.args = args
+        self.args: Any = args
 
         cli.exit_if_empty(args.text, parser)
         result = self.run(