Fix some types

wonderwords/_trie.py

@@ -53,7 +53,7 @@ def get_words_that_start_with(self, characters: str) -> Set[str]:
 
     def get_words_from_branch(self, branch: TrieNode, word_fragment: str) -> Set[str]:
         """Get all words that start with ``word_fragment`` starting from the node ``branch``."""
-        words = set()
+        words: Set[str] = set()
 
         if branch.end_of_word:
             words.add(word_fragment)

wonderwords/cmdline.py

@@ -1,3 +1,5 @@
+from typing import List
+
 from rich import print
 from rich.markdown import Markdown
 
@@ -26,20 +28,20 @@ def display_word(word: str) -> None:
     print(f"[cyan]{word}[/cyan]")
 
 
-def display_list(words: list, delimiter=",") -> None:
+def display_list(words: List[str], delimiter: str = ",") -> None:
     delimiter_colorized = f"[grey50]{delimiter}[/grey50]"
     print(delimiter_colorized.join([f"[cyan]{word}[/cyan]" for word in words]))
 
 
-def display_word_not_found(one_word=True) -> None:
+def display_word_not_found(one_word: bool = True) -> None:
     print(
         f"[red]:exclamation: No word{'' if one_word else 's'} matching the criteria specified could be found.[/red]"
     )
 
 
 def display_not_enough_words() -> None:
     print(
-        f"[red]:exclamation: Couldn't find enough words matching the criteria specified.[/red]"
+        "[red]:exclamation: Couldn't find enough words matching the criteria specified.[/red]"
     )
 
 

wonderwords/cmdline_parser.py

@@ -169,7 +169,7 @@ def get_mode(arguments):
         return None
 
 
-def run_wonderwords(mode, arguments):  # noqa: C901
+def run_wonderwords(mode: str, arguments):  # noqa: C901
     if mode == "version":
         cmdline.display_version()
         return 0

wonderwords/random_word.py

@@ -92,7 +92,7 @@ def _load_default_categories(
     default_categories: Type[Defaults],
 ) -> Dict[Defaults, WordList]:
     """Load all the default word lists"""
-    out = {}
+    out: Dict[Defaults, WordList] = {}
     for category in default_categories:
         out[category] = _get_words_from_text_file(category.value)
     return out
@@ -169,7 +169,7 @@ class RandomWord:
     """
 
     def __init__(
-        self, enhanced_prefixes: bool = True, rng=None, **kwargs: Union[WordList, Defaults]
+        self, enhanced_prefixes: bool = True, rng: Optional[Random] = None, **kwargs: Union[WordList, Defaults]
     ):
         # A dictionary where lists of words organized into named categories
         self._categories: Dict[str, WordList]
@@ -280,7 +280,7 @@ def filter(  # noqa: C901
         # are done at once since categories are specifically ordered
         # in order to make filtering by length an efficient process.
         # See issue #14 for details.
-        words = set()
+        words: Set[str] = set()
 
         for category in include_categories:
             try:
@@ -426,7 +426,7 @@ def random_words(
                     f"{str(amount)} word(s)"
                 )
 
-        words = []
+        words: WordList = []
         for _ in range(amount):
             new_word = self._generator.choice(choose_from)
             choose_from.remove(new_word)
@@ -532,7 +532,7 @@ def _get_word_lists_by_category(
         self, custom_categories: Dict[str, Any]
     ) -> Dict[str, WordList]:
         """Add custom categories of words"""
-        out = {}
+        out: dict[str, WordList] = {}
         for name, words in custom_categories.items():
             if isinstance(words, Defaults):
                 word_list = _DEFAULT_CATEGORIES[words]
@@ -591,7 +591,7 @@ def _perform_long_operations(
         self, words: Set[str], long_operations: Dict[str, Any]
     ) -> Set[str]:
         """Return a set of words that do not meet the criteria specified by the long operations."""
-        remove_words = set()
+        remove_words: Set[str] = set()
         for word in words:
             if "regex" in long_operations:
                 if not long_operations["regex"].fullmatch(word):