Merge pull request #39 from plasma-umass/typerwriter_annotations

type annotations

src/chatdbg/assistant/assistant.py

@@ -22,7 +22,7 @@ def __init__(self, *args: object) -> None:
         super().__init__(*args)
 
 
-def remove_non_printable_chars(s):
+def remove_non_printable_chars(s: str) -> str:
     printable_chars = set(string.printable)
     filtered_string = "".join(filter(lambda x: x in printable_chars, s))
     return filtered_string

src/chatdbg/pdb_util/locals.py

@@ -82,7 +82,7 @@ def _repr_if_defined(obj: Any) -> bool:
     return result
 
 
-def _format_limited(value: Any, limit: int = 10, depth: int = 3) -> str:
+def _format_limited(value: Union[int, np.ndarray], limit: int = 10, depth: int = 3) -> str:
     def format_tuple(t, depth):
         return tuple([helper(x, depth) for x in t])
 

src/chatdbg/util/config.py

@@ -17,8 +17,8 @@
 
 
 def _chatdbg_get_env(
-    option_name: str, default_value: Union[int, str, bool]
-) -> Union[int, str, bool]:
+    option_name: str, default_value: Union[int, bool, str]
+) -> Union[int, bool, str]:
     env_name = "CHATDBG_" + option_name.upper()
     v = os.getenv(env_name, str(default_value))
     if type(default_value) == int:

src/chatdbg/util/prompts.py

@@ -2,9 +2,10 @@
 import os
 from chatdbg.util.config import chatdbg_config
 from .text import truncate_proportionally
+from types import NoneType
+from typing import Any, Callable, List, Union
 
-
-def _wrap_it(before, x, after="", maxlen=2048):
+def _wrap_it(before: str, x: Union[str, NoneType], after: str="", maxlen: int=2048) -> str:
     if x:
         x = truncate_proportionally(x, maxlen, 0.5)
         before = before + ":\n" if before else ""
@@ -14,18 +15,18 @@ def _wrap_it(before, x, after="", maxlen=2048):
         return ""
 
 
-def _concat_prompt(*args):
+def _concat_prompt(*args) -> str:
     args = [a for a in args if a]
     return "\n".join(args)
 
 
-def _user_text_it(user_text):
+def _user_text_it(user_text: str) -> str:
     return user_text if user_text else "What's the bug? Give me a fix."
 
 
 def build_initial_prompt(
-    stack, error, details, command_line, inputs, history, extra="", user_text=""
-):
+    stack: str, error: str, details: str, command_line: str, inputs: str, history: str, extra: NoneType="", user_text: str=""
+) -> str:
     return _concat_prompt(
         _wrap_it("The program has this stack trace", stack),
         _wrap_it("The program encountered the following error", error, details),
@@ -37,15 +38,15 @@ def build_initial_prompt(
     )
 
 
-def build_followup_prompt(history, extra, user_text):
+def build_followup_prompt(history: str, extra: str, user_text: str) -> str:
     return _concat_prompt(
         _wrap_it("This is the history of some debugger commands I ran", history),
         _wrap_it("", extra),
         _user_text_it(user_text),
     )
 
 
-def initial_instructions(functions):
+def initial_instructions(functions: List[Callable[[Any], Any]]) -> str:
     if chatdbg_config.instructions == "":
         file_path = os.path.join(os.path.dirname(__file__), "instructions.txt")
     else:

src/chatdbg/util/text.py

@@ -1,6 +1,6 @@
 import re
 import textwrap
-
+from typing import Union
 
 def make_arrow(pad):
     """generate the leading arrow in front of traceback or debugger"""
@@ -16,7 +16,7 @@ def strip_ansi(s: str) -> str:
     return ansi_escape.sub("", s)
 
 
-def truncate_proportionally(text, maxlen=32000, top_proportion=0.5):
+def truncate_proportionally(text: str, maxlen: int=32000, top_proportion: Union[int, float]=0.5) -> str:
     """Omit part of a string if needed to make it fit in a maximum length."""
     if len(text) > maxlen:
         pre = max(0, int((maxlen - 5) * top_proportion))
@@ -25,7 +25,7 @@ def truncate_proportionally(text, maxlen=32000, top_proportion=0.5):
     return text
 
 
-def wrap_long_lines(text, width=80, subsequent_indent="    "):
+def wrap_long_lines(text: str, width: int=80, subsequent_indent: str="    ") -> str:
     wrapped_lines = []
     for line in text.split("\n"):
         if len(line) > width:
@@ -37,5 +37,5 @@ def wrap_long_lines(text, width=80, subsequent_indent="    "):
     return "\n".join(wrapped_lines)
 
 
-def fill_to_width(text, width=80):
+def fill_to_width(text: str, width: int=80) -> str:
     return "\n".join([line.ljust(width) for line in text.split("\n")])

src/chatdbg/util/trim.py

@@ -1,5 +1,6 @@
 import copy
 import warnings
+from typing import Dict, List, Union
 
 with warnings.catch_warnings():
     warnings.simplefilter("ignore")
@@ -65,10 +66,10 @@ def chunkify(messages, model):
 
 
 def trim_messages(
-    messages,
-    model,
+    messages: List[Dict[str, str]],
+    model: str,
     trim_ratio: float = 0.75,
-):
+) -> list:
     """
     Strategy:
     - chunk messages: