Type hints for completion_refresher and sqlexecute

Author: amjith

litecli/completion_refresher.py

@@ -1,6 +1,10 @@
 # mypy: ignore-errors
 
+from __future__ import annotations
+
 import threading
+from typing import Callable, Dict, List, Optional, Tuple
+
 from .packages.special.main import COMMANDS
 from collections import OrderedDict
 
@@ -9,13 +13,18 @@
 
 
 class CompletionRefresher(object):
-    refreshers = OrderedDict()
+    refreshers: Dict[str, Callable] = OrderedDict()
 
-    def __init__(self):
-        self._completer_thread = None
+    def __init__(self) -> None:
+        self._completer_thread: Optional[threading.Thread] = None
         self._restart_refresh = threading.Event()
 
-    def refresh(self, executor, callbacks, completer_options=None):
+    def refresh(
+        self,
+        executor: SQLExecute,
+        callbacks: Callable | List[Callable],
+        completer_options: Optional[dict] = None,
+    ) -> List[Tuple]:
         """Creates a SQLCompleter object and populates it with the relevant
         completion suggestions in a background thread.
 
@@ -34,7 +43,7 @@ def refresh(self, executor, callbacks, completer_options=None):
             self._restart_refresh.set()
             return [(None, None, None, "Auto-completion refresh restarted.")]
         else:
-            if executor.dbname == ":memory:":
+            if executor.dbname == ":memory":
                 # if DB is memory, needed to use same connection
                 # So can't use same connection with different thread
                 self._bg_refresh(executor, callbacks, completer_options)
@@ -46,19 +55,17 @@ def refresh(self, executor, callbacks, completer_options=None):
                 )
                 self._completer_thread.daemon = True
                 self._completer_thread.start()
-                return [
-                    (
-                        None,
-                        None,
-                        None,
-                        "Auto-completion refresh started in the background.",
-                    )
-                ]
-
-    def is_refreshing(self):
-        return self._completer_thread and self._completer_thread.is_alive()
-
-    def _bg_refresh(self, sqlexecute, callbacks, completer_options):
+                return [(None, None, None, "Auto-completion refresh started in the background.")]
+
+    def is_refreshing(self) -> bool:
+        return bool(self._completer_thread and self._completer_thread.is_alive())
+
+    def _bg_refresh(
+        self,
+        sqlexecute: SQLExecute,
+        callbacks: Callable | List[Callable],
+        completer_options: dict,
+    ) -> None:
         completer = SQLCompleter(**completer_options)
 
         e = sqlexecute
@@ -92,41 +99,42 @@ def _bg_refresh(self, sqlexecute, callbacks, completer_options):
             callback(completer)
 
 
-def refresher(name, refreshers=CompletionRefresher.refreshers):
+def refresher(name: str, refreshers: Dict[str, Callable] = CompletionRefresher.refreshers):
     """Decorator to add the decorated function to the dictionary of
     refreshers. Any function decorated with a @refresher will be executed as
     part of the completion refresh routine."""
 
-    def wrapper(wrapped):
+    def wrapper(wrapped: Callable) -> Callable:
         refreshers[name] = wrapped
         return wrapped
 
     return wrapper
 
 
 @refresher("databases")
-def refresh_databases(completer, executor):
+def refresh_databases(completer: SQLCompleter, executor: SQLExecute) -> None:
     completer.extend_database_names(executor.databases())
 
 
 @refresher("schemata")
-def refresh_schemata(completer, executor):
+def refresh_schemata(completer: SQLCompleter, executor: SQLExecute) -> None:
     # name of the current database.
     completer.extend_schemata(executor.dbname)
     completer.set_dbname(executor.dbname)
 
 
 @refresher("tables")
-def refresh_tables(completer, executor):
-    completer.extend_relations(executor.tables(), kind="tables")
-    completer.extend_columns(executor.table_columns(), kind="tables")
+def refresh_tables(completer: SQLCompleter, executor: SQLExecute) -> None:
+    table_cols = list(executor.table_columns())
+    completer.extend_relations(table_cols, kind="tables")
+    completer.extend_columns(table_cols, kind="tables")
 
 
 @refresher("functions")
-def refresh_functions(completer, executor):
+def refresh_functions(completer: SQLCompleter, executor: SQLExecute) -> None:
     completer.extend_functions(executor.functions())
 
 
 @refresher("special_commands")
-def refresh_special(completer, executor):
-    completer.extend_special_commands(COMMANDS.keys())
+def refresh_special(completer: SQLCompleter, executor: SQLExecute) -> None:
+    completer.extend_special_commands(list(COMMANDS.keys()))

litecli/sqlexecute.py

@@ -1,6 +1,9 @@
 # mypy: ignore-errors
 
+from __future__ import annotations
+
 import logging
+from typing import Any, Generator, Iterable, Optional, Tuple
 
 from contextlib import closing
 
@@ -57,7 +60,7 @@ class SQLExecute(object):
     functions_query = '''SELECT ROUTINE_NAME FROM INFORMATION_SCHEMA.ROUTINES
     WHERE ROUTINE_TYPE="FUNCTION" AND ROUTINE_SCHEMA = "%s"'''
 
-    def __init__(self, database):
+    def __init__(self, database: Optional[str]):
         self.dbname = database
         self._server_type = None
         self.conn = None
@@ -66,7 +69,7 @@ def __init__(self, database):
             return
         self.connect()
 
-    def connect(self, database=None):
+    def connect(self, database: Optional[str] = None) -> None:
         db = database or self.dbname
         _logger.debug("Connection DB Params: \n\tdatabase: %r", db)
 
@@ -85,7 +88,7 @@ def connect(self, database=None):
         # successful connection.
         self.dbname = db
 
-    def run(self, statement):
+    def run(self, statement: str) -> Iterable[Tuple]:
         """Execute the sql in the database and return the results. The results
         are a list of tuples. Each tuple has 4 values
         (title, rows, headers, status).
@@ -140,7 +143,7 @@ def run(self, statement):
                     cur.execute(sql)
                     yield self.get_result(cur)
 
-    def get_result(self, cursor):
+    def get_result(self, cursor: Any) -> Tuple[Optional[str], Optional[list], Optional[list], str]:
         """Get the current result's data from the cursor."""
         title = headers = None
 
@@ -164,7 +167,7 @@ def get_result(self, cursor):
 
         return (title, cursor, headers, status)
 
-    def tables(self):
+    def tables(self) -> Generator[Tuple[str], None, None]:
         """Yields table names"""
 
         with closing(self.conn.cursor()) as cur:
@@ -173,15 +176,15 @@ def tables(self):
             for row in cur:
                 yield row
 
-    def table_columns(self):
+    def table_columns(self) -> Generator[Tuple[str, str], None, None]:
         """Yields column names"""
         with closing(self.conn.cursor()) as cur:
             _logger.debug("Columns Query. sql: %r", self.table_columns_query)
             cur.execute(self.table_columns_query)
             for row in cur:
                 yield row
 
-    def databases(self):
+    def databases(self) -> Iterable[str]:
         if not self.conn:
             return
 
@@ -190,7 +193,7 @@ def databases(self):
             for row in cur.execute(self.databases_query):
                 yield row[1]
 
-    def functions(self):
+    def functions(self) -> Iterable[Tuple]:
         """Yields tuples of (schema_name, function_name)"""
 
         with closing(self.conn.cursor()) as cur:
@@ -199,7 +202,7 @@ def functions(self):
             for row in cur:
                 yield row
 
-    def show_candidates(self):
+    def show_candidates(self) -> Iterable[Any]:
         with closing(self.conn.cursor()) as cur:
             _logger.debug("Show Query. sql: %r", self.show_candidates_query)
             try:
@@ -211,6 +214,6 @@ def show_candidates(self):
                 for row in cur:
                     yield (row[0].split(None, 1)[-1],)
 
-    def server_type(self):
+    def server_type(self) -> Tuple[str, str]:
         self._server_type = ("sqlite3", "3")
         return self._server_type