refine typehints in special/llm.py

 * "Optional" can be replaced by "| None" in modern Pythons
 * "Tuple" can be replaced by lowercase "tuple" in modern Pythons
 * variable "cur" has a type of pymysql.cursors.Cursor
 * typehint contextlib variable "redirect"
 * change exit code to "int(e.code or 0)", a semantic change
 * split build_command_tree() using an inner _build_command_tree() to
   simplify the return type
 * don't pass mutable COMMAND_TREE as an argument default
 * typehint almost all function arguments
 * typehint all return types
 * remove a needless "return" statement
 * reformat many parameter lists as vertical

Author: rolandwalker

changelog.md

@@ -5,6 +5,7 @@ Internal
 --------
 * Add mypy to Pull Request template.
 * Enable flake8-bugbear lint rules.
+* Improve typehints on LLM driver.
 
 
 1.40.0 (2025/10/14)

mycli/packages/special/llm.py

@@ -7,11 +7,12 @@
 import shlex
 import sys
 from time import time
-from typing import Optional, Tuple
+from typing import Any
 
 import click
 import llm
 from llm.cli import cli
+from pymysql.cursors import Cursor
 
 from mycli.packages.special.main import Verbosity, parse_special_command
 
@@ -22,15 +23,22 @@
 LLM_TEMPLATE_NAME = "mycli-llm-template"
 
 
-def run_external_cmd(cmd, *args, capture_output=False, restart_cli=False, raise_exception=True):
+def run_external_cmd(
+    cmd: str,
+    *args,
+    capture_output=False,
+    restart_cli=False,
+    raise_exception=True,
+) -> tuple[int, str]:
     original_exe = sys.executable
     original_args = sys.argv
     try:
         sys.argv = [cmd] + list(args)
         code = 0
         if capture_output:
             buffer = io.StringIO()
-            redirect = contextlib.ExitStack()
+            redirect: contextlib.ExitStack[bool | None] | contextlib.nullcontext[None] = contextlib.ExitStack()
+            assert isinstance(redirect, contextlib.ExitStack)
             redirect.enter_context(contextlib.redirect_stdout(buffer))
             redirect.enter_context(contextlib.redirect_stderr(buffer))
         else:
@@ -39,7 +47,7 @@ def run_external_cmd(cmd, *args, capture_output=False, restart_cli=False, raise_
             try:
                 run_module(cmd, run_name="__main__")
             except SystemExit as e:
-                code = e.code
+                code = int(e.code or 0)
                 if code != 0 and raise_exception:
                     if capture_output:
                         raise RuntimeError(buffer.getvalue())
@@ -62,24 +70,33 @@ def run_external_cmd(cmd, *args, capture_output=False, restart_cli=False, raise_
         sys.argv = original_args
 
 
-def build_command_tree(cmd):
-    tree = {}
+def _build_command_tree(cmd) -> dict[str, Any] | None:
+    tree: dict[str, Any] | None = {}
+    assert isinstance(tree, dict)
     if isinstance(cmd, click.Group):
         for name, subcmd in cmd.commands.items():
             if cmd.name == "models" and name == "default":
                 tree[name] = MODELS
             else:
-                tree[name] = build_command_tree(subcmd)
+                tree[name] = _build_command_tree(subcmd)
     else:
         tree = None
     return tree
 
 
+def build_command_tree(cmd) -> dict[str, Any]:
+    return _build_command_tree(cmd) or {}
+
+
 # Generate the command tree for autocompletion
 COMMAND_TREE = build_command_tree(cli) if cli else {}
 
 
-def get_completions(tokens, tree=COMMAND_TREE):
+def get_completions(
+    tokens: list[str],
+    tree: dict[str, Any] | None = None,
+) -> list[str]:
+    tree = tree or COMMAND_TREE
     for token in tokens:
         if token.startswith("-"):
             continue
@@ -150,16 +167,15 @@ def __init__(self, results=None):
 """
 
 
-def ensure_mycli_template(replace=False):
+def ensure_mycli_template(replace: bool = False) -> None:
     if not replace:
         code, _ = run_external_cmd("llm", "templates", "show", LLM_TEMPLATE_NAME, capture_output=True, raise_exception=False)
         if code == 0:
             return
     run_external_cmd("llm", PROMPT, "--save", LLM_TEMPLATE_NAME)
-    return
 
 
-def handle_llm(text, cur) -> Tuple[str, Optional[str], float]:
+def handle_llm(text: str, cur: Cursor) -> tuple[str, str | None, float]:
     _, verbosity, arg = parse_special_command(text)
     if not arg.strip():
         output = [(None, None, None, USAGE)]
@@ -214,12 +230,15 @@ def handle_llm(text, cur) -> Tuple[str, Optional[str], float]:
         raise RuntimeError(e)
 
 
-def is_llm_command(command) -> bool:
+def is_llm_command(command: str) -> bool:
     cmd, _, _ = parse_special_command(command)
     return cmd in ("\\llm", "\\ai")
 
 
-def sql_using_llm(cur, question=None) -> Tuple[str, Optional[str]]:
+def sql_using_llm(
+    cur: Cursor | None,
+    question: str | None = None,
+) -> tuple[str, str | None]:
     if cur is None:
         raise RuntimeError("Connect to a database and try again.")
     schema_query = """