Type hint main.py and introduce DBCursor protocol.

Author: amjith

litecli/main.py

@@ -1,6 +1,6 @@
 # mypy: ignore-errors
 
-from __future__ import print_function, unicode_literals
+from __future__ import annotations
 
 import itertools
 import logging
@@ -19,6 +19,7 @@
 except ImportError:
     from sqlite3 import OperationalError, sqlite_version
 from time import time
+from typing import Any, Iterable, Optional, Dict
 
 import click
 import sqlparse
@@ -36,6 +37,7 @@
 )
 from prompt_toolkit.lexers import PygmentsLexer
 from prompt_toolkit.shortcuts import CompleteStyle, PromptSession
+from prompt_toolkit.completion import Completion
 
 from .__init__ import __version__
 from .clibuffer import cli_is_multiline
@@ -66,13 +68,13 @@ class LiteCli(object):
 
     def __init__(
         self,
-        sqlexecute=None,
-        prompt=None,
-        logfile=None,
-        auto_vertical_output=False,
-        warn=None,
-        liteclirc=None,
-    ):
+        sqlexecute: Optional[SQLExecute] = None,
+        prompt: Optional[str] = None,
+        logfile: Optional[Any] = None,
+        auto_vertical_output: bool = False,
+        warn: Optional[bool] = None,
+        liteclirc: Optional[str] = None,
+    ) -> None:
         self.sqlexecute = sqlexecute
         self.logfile = logfile
 
@@ -139,7 +141,7 @@ def __init__(
 
         self.prompt_app = None
 
-    def register_special_commands(self):
+    def register_special_commands(self) -> None:
         special.register_special_command(
             self.change_db,
             ".open",
@@ -180,7 +182,7 @@ def register_special_commands(self):
             case_sensitive=True,
         )
 
-    def change_table_format(self, arg, **_):
+    def change_table_format(self, arg: str, **_: Any):
         try:
             self.formatter.format_name = arg
             yield (None, None, None, "Changed table format to {}".format(arg))
@@ -190,7 +192,7 @@ def change_table_format(self, arg, **_):
                 msg += "\n\t{}".format(table_type)
             yield (None, None, None, msg)
 
-    def change_db(self, arg, **_):
+    def change_db(self, arg: Optional[str], **_: Any):
         if arg is None:
             self.sqlexecute.connect()
         else:
@@ -204,7 +206,7 @@ def change_db(self, arg, **_):
             'You are now connected to database "%s"' % (self.sqlexecute.dbname),
         )
 
-    def execute_from_file(self, arg, **_):
+    def execute_from_file(self, arg: Optional[str], **_: Any):
         if not arg:
             message = "Missing required argument, filename."
             return [(None, None, None, message)]
@@ -220,7 +222,7 @@ def execute_from_file(self, arg, **_):
 
         return self.sqlexecute.run(query)
 
-    def change_prompt_format(self, arg, **_):
+    def change_prompt_format(self, arg: Optional[str], **_: Any):
         """
         Change the prompt format.
         """
@@ -231,7 +233,7 @@ def change_prompt_format(self, arg, **_):
         self.prompt_format = self.get_prompt(arg)
         return [(None, None, None, "Changed prompt format to %s" % arg)]
 
-    def initialize_logging(self):
+    def initialize_logging(self) -> None:
         log_file = self.config["main"]["log_file"]
         if log_file == "default":
             log_file = config_location() + "log"
@@ -279,7 +281,7 @@ def initialize_logging(self):
         root_logger.debug("Initializing litecli logging.")
         root_logger.debug("Log file %r.", log_file)
 
-    def read_my_cnf_files(self, keys):
+    def read_my_cnf_files(self, keys: Iterable[str]) -> Dict[str, Optional[str]]:
         """
         Reads a list of config files and merges them. The last one will win.
         :param files: list of files to read
@@ -299,7 +301,7 @@ def get(key):
 
         return {x: get(x) for x in keys}
 
-    def connect(self, database=""):
+    def connect(self, database: str = "") -> None:
         cnf = {"database": None}
 
         cnf = self.read_my_cnf_files(cnf.keys())
@@ -321,7 +323,7 @@ def _connect():
             self.echo(str(e), err=True, fg="red")
             exit(1)
 
-    def handle_editor_command(self, text):
+    def handle_editor_command(self, text: str) -> str:
         R"""Editor command is any query that is prefixed or suffixed by a '\e'.
         The reason for a while loop is because a user might edit a query
         multiple times. For eg:
@@ -351,7 +353,7 @@ def handle_editor_command(self, text):
             continue
         return text
 
-    def run_cli(self):
+    def run_cli(self) -> None:
         iterations = 0
         sqlexecute = self.sqlexecute
         logger = self.logger
@@ -645,12 +647,12 @@ def startup_commands():
             if not self.less_chatty:
                 self.echo("Goodbye!")
 
-    def log_output(self, output):
+    def log_output(self, output: str) -> None:
         """Log the output in the audit log, if it's enabled."""
         if self.logfile:
             click.echo(output, file=self.logfile)
 
-    def echo(self, s, **kwargs):
+    def echo(self, s: str, **kwargs: Any) -> None:
         """Print a message to stdout.
 
         The message will be logged in the audit log, if enabled.
@@ -661,7 +663,7 @@ def echo(self, s, **kwargs):
         self.log_output(s)
         click.secho(s, **kwargs)
 
-    def get_output_margin(self, status=None):
+    def get_output_margin(self, status: Optional[str] = None) -> int:
         """Get the output margin (number of rows for the prompt, footer and
         timing message."""
         margin = self.get_reserved_space() + self.get_prompt(self.prompt_format).count("\n") + 2
@@ -670,7 +672,7 @@ def get_output_margin(self, status=None):
 
         return margin
 
-    def output(self, output, status=None):
+    def output(self, output: Iterable[str], status: Optional[str] = None) -> None:
         """Output text to stdout or a pager command.
 
         The status text is not outputted to pager or files.
@@ -723,7 +725,7 @@ def output(self, output, status=None):
             self.log_output(status)
             click.secho(status)
 
-    def configure_pager(self):
+    def configure_pager(self) -> None:
         # Provide sane defaults for less if they are empty.
         if not os.environ.get("LESS"):
             os.environ["LESS"] = "-RXF"
@@ -738,7 +740,7 @@ def configure_pager(self):
         if cnf["skip-pager"] or not self.config["main"].as_bool("enable_pager"):
             special.disable_pager()
 
-    def refresh_completions(self, reset=False):
+    def refresh_completions(self, reset: bool = False):
         if reset:
             with self._completer_lock:
                 self.completer.reset_completions()
@@ -753,7 +755,7 @@ def refresh_completions(self, reset=False):
 
         return [(None, None, None, "Auto-completion refresh started in the background.")]
 
-    def _on_completions_refreshed(self, new_completer):
+    def _on_completions_refreshed(self, new_completer: SQLCompleter) -> None:
         """Swap the completer object in cli with the newly created completer."""
         with self._completer_lock:
             self.completer = new_completer
@@ -763,11 +765,11 @@ def _on_completions_refreshed(self, new_completer):
             # "Refreshing completions..." indicator
             self.prompt_app.app.invalidate()
 
-    def get_completions(self, text, cursor_positition):
+    def get_completions(self, text: str, cursor_positition: int) -> Iterable[Completion]:
         with self._completer_lock:
             return self.completer.get_completions(Document(text=text, cursor_position=cursor_positition), None)
 
-    def get_prompt(self, string):
+    def get_prompt(self, string: str) -> str:
         self.logger.debug("Getting prompt %r", string)
         sqlexecute = self.sqlexecute
         now = datetime.now()
@@ -795,7 +797,7 @@ def replacer(match):
         # Perform the substitution
         return pattern.sub(replacer, string)
 
-    def run_query(self, query, new_line=True):
+    def run_query(self, query: str, new_line: bool = True):
         """Runs *query*."""
         results = self.sqlexecute.run(query)
         for result in results:

litecli/packages/special/dbcommands.py

@@ -8,7 +8,7 @@
 import sys
 import platform
 import shlex
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional, Tuple, Protocol, Sequence
 
 
 from litecli import __version__
@@ -18,6 +18,16 @@
 log = logging.getLogger(__name__)
 
 
+class DBCursor(Protocol):
+    description: Optional[Sequence[Sequence[Any]]]
+
+    def execute(self, sql: str, params: Any = ...) -> Any: ...
+
+    def fetchall(self) -> List[Tuple[Any, ...]]: ...
+
+    def fetchone(self) -> Optional[Tuple[Any, ...]]: ...
+
+
 @special_command(
     ".tables",
     "\\dt",
@@ -27,7 +37,7 @@
     aliases=("\\dt",),
 )
 def list_tables(
-    cur: Any,
+    cur: DBCursor,
     arg: Optional[str] = None,
     arg_type: int = PARSED_QUERY,
     verbose: bool = False,
@@ -74,7 +84,7 @@ def list_tables(
     aliases=("\\dv",),
 )
 def list_views(
-    cur: Any,
+    cur: DBCursor,
     arg: Optional[str] = None,
     arg_type: int = PARSED_QUERY,
     verbose: bool = False,
@@ -111,7 +121,7 @@ def list_views(
     arg_type=PARSED_QUERY,
     case_sensitive=True,
 )
-def show_schema(cur: Any, arg: Optional[str] = None, **_: Any) -> List[Tuple]:
+def show_schema(cur: DBCursor, arg: Optional[str] = None, **_: Any) -> List[Tuple]:
     if arg:
         args = (arg,)
         query = """
@@ -147,7 +157,7 @@ def show_schema(cur: Any, arg: Optional[str] = None, **_: Any) -> List[Tuple]:
     case_sensitive=True,
     aliases=("\\l",),
 )
-def list_databases(cur: Any, **_: Any) -> List[Tuple]:
+def list_databases(cur: DBCursor, **_: Any) -> List[Tuple]:
     query = "PRAGMA database_list"
     log.debug(query)
     cur.execute(query)
@@ -167,7 +177,7 @@ def list_databases(cur: Any, **_: Any) -> List[Tuple]:
     aliases=("\\di",),
 )
 def list_indexes(
-    cur: Any,
+    cur: DBCursor,
     arg: Optional[str] = None,
     arg_type: int = PARSED_QUERY,
     verbose: bool = False,
@@ -206,7 +216,7 @@ def list_indexes(
     aliases=("\\s",),
     case_sensitive=True,
 )
-def status(cur: Any, **_: Any) -> List[Tuple]:
+def status(cur: DBCursor, **_: Any) -> List[Tuple]:
     # Create output buffers.
     footer = []
     footer.append("--------------")
@@ -267,7 +277,7 @@ def load_extension(cur, arg, **_):
     case_sensitive=True,
     aliases=("\\d", "desc"),
 )
-def describe(cur: Any, arg: Optional[str], **_: Any) -> List[Tuple]:
+def describe(cur: DBCursor, arg: Optional[str], **_: Any) -> List[Tuple]:
     if arg:
         query = """
             PRAGMA table_info({})
@@ -294,7 +304,7 @@ def describe(cur: Any, arg: Optional[str], **_: Any) -> List[Tuple]:
     arg_type=PARSED_QUERY,
     case_sensitive=True,
 )
-def import_file(cur: Any, arg: Optional[str] = None, **_: Any) -> List[Tuple]:
+def import_file(cur: DBCursor, arg: Optional[str] = None, **_: Any) -> List[Tuple]:
     def split(s):
         # this is a modification of shlex.split function, just to make it support '`',
         # because table name might contain '`' character.

litecli/packages/special/iocommands.py

@@ -10,7 +10,7 @@
 import subprocess
 from io import open
 from time import sleep
-from typing import Any, Dict, Generator, List, Optional, Tuple
+from typing import Any, Dict, Generator, List, Optional, Tuple, TextIO
 
 import click
 import sqlparse
@@ -24,10 +24,10 @@
 
 use_expanded_output: bool = False
 PAGER_ENABLED: bool = True
-tee_file: Any = None
-once_file: Any = None
-written_to_once_file: Optional[bool] = None
-pipe_once_process: Any = None
+tee_file: Optional[TextIO] = None
+once_file: Optional[TextIO] = None
+written_to_once_file: bool = False
+pipe_once_process: Optional[subprocess.Popen[str]] = None
 written_to_pipe_once_process: bool = False
 favoritequeries: FavoriteQueries = FavoriteQueries(ConfigObj())
 

litecli/packages/special/llm.py

@@ -16,7 +16,7 @@
 import sys
 from runpy import run_module
 from time import time
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Protocol, Sequence
 
 import click
 import llm
@@ -219,7 +219,18 @@ def ensure_litecli_template(replace: bool = False) -> None:
 
 
 @export
-def handle_llm(text: str, cur: Any) -> Tuple[str, Optional[str], float]:
+class DBCursor(Protocol):
+    description: Optional[Sequence[Sequence[Any]]]
+
+    def execute(self, sql: str, params: Any = ...) -> Any: ...
+
+    def fetchall(self) -> List[Tuple[Any, ...]]: ...
+
+    def fetchone(self) -> Optional[Tuple[Any, ...]]: ...
+
+
+@export
+def handle_llm(text: str, cur: DBCursor) -> Tuple[str, Optional[str], float]:
     """This function handles the special command `\\llm`.
 
     If it deals with a question that results in a SQL query then it will return
@@ -323,7 +334,7 @@ def is_llm_command(command: str) -> bool:
 
 @export
 def sql_using_llm(
-    cur: Any,
+    cur: DBCursor,
     question: Optional[str] = None,
     verbose: bool = False,
 ) -> Tuple[str, Optional[str], Optional[str]]: