From db0487548fcd6e17ab46d4af18ee564f49db6475 Mon Sep 17 00:00:00 2001
From: Nicolas van Kempen <nvankemp@gmail.com>
Date: Wed, 22 May 2024 13:33:42 -0400
Subject: [PATCH] A few more type annotations

---
 src/cwhy/__main__.py | 6 +++---
 src/cwhy/cwhy.py     | 9 ++++++---
 tests/prepare.py     | 2 +-
 tests/regression.py  | 4 ++--
 4 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/cwhy/__main__.py b/src/cwhy/__main__.py
index 3c45bd5..0e4fd1b 100755
--- a/src/cwhy/__main__.py
+++ b/src/cwhy/__main__.py
@@ -60,16 +60,16 @@ def _print_message(self, message: Optional[str], file: Any = None) -> None:
 
 class CWhyArgumentFormatter(argparse.HelpFormatter):
     # RawDescriptionHelpFormatter.
-    def _fill_text(self, text, width, indent):
+    def _fill_text(self, text, width, indent):  # type: ignore
         return "".join(indent + line for line in text.splitlines(keepends=True))
 
     # RawTextHelpFormatter.
-    def _split_lines(self, text, width):
+    def _split_lines(self, text, width):  # type: ignore
         return text.splitlines()
 
     # ArgumentDefaultsHelpFormatter.
     # Ignore if help message is multiline.
-    def _get_help_string(self, action):
+    def _get_help_string(self, action):  # type: ignore
         help = action.help
         if "\n" not in help and "%(default)" not in action.help:
             if action.default is not argparse.SUPPRESS:
diff --git a/src/cwhy/cwhy.py b/src/cwhy/cwhy.py
index 3457841..944621b 100755
--- a/src/cwhy/cwhy.py
+++ b/src/cwhy/cwhy.py
@@ -2,6 +2,7 @@
 import os
 import subprocess
 import sys
+from typing import Any
 import warnings
 
 with warnings.catch_warnings():
@@ -56,7 +57,7 @@ def print_key_info() -> None:
     sys.exit(1)
 
 
-def complete(args, user_prompt, **kwargs):
+def complete(args: argparse.Namespace, user_prompt: str, **kwargs: Any):
     try:
         completion = litellm.completion(
             model=args.llm,
@@ -198,13 +199,15 @@ def main(args: argparse.Namespace) -> None:
     sys.exit(process.returncode)
 
 
-def evaluate_text_prompt(args, prompt, wrap=True, **kwargs):
+def evaluate_text_prompt(
+    args: argparse.Namespace, prompt: str, wrap: bool = True, **kwargs: Any
+) -> str:
     completion = complete(args, prompt, **kwargs)
 
     msg = f"Analysis from {args.llm}:"
     dprint(msg)
     dprint("-" * len(msg))
-    text = completion.choices[0].message.content
+    text: str = completion.choices[0].message.content
 
     if wrap:
         text = llm_utils.word_wrap_except_code_blocks(text)
diff --git a/tests/prepare.py b/tests/prepare.py
index e0a5c75..fb2fc62 100644
--- a/tests/prepare.py
+++ b/tests/prepare.py
@@ -7,7 +7,7 @@
 ROOT = os.path.dirname(os.path.abspath(__file__))
 
 
-def prepare(language) -> None:
+def prepare(language: str) -> None:
     try:
         with open(os.path.join(ROOT, language, "manifest.yml"), "r") as stream:
             manifest = yaml.load(stream, yaml.Loader)
diff --git a/tests/regression.py b/tests/regression.py
index 3e434e0..06244b9 100644
--- a/tests/regression.py
+++ b/tests/regression.py
@@ -10,7 +10,7 @@
 ROOT = os.path.dirname(os.path.abspath(__file__))
 
 
-def get_cwhy_prompt(invocation) -> str:
+def get_cwhy_prompt(invocation: str) -> str:
     process = subprocess.run(
         f"cwhy --show-prompt --- {invocation}",
         text=True,
@@ -28,7 +28,7 @@ def get_cwhy_prompt(invocation) -> str:
     return stdout
 
 
-def main(args) -> None:
+def main(args: argparse.Namespace) -> None:
     prepare.clean()
     prepare.prepare_all()