Type hints for sqlcompleter.

Author: amjith

litecli/sqlcompleter.py

@@ -1,12 +1,16 @@
+"""SQL completer with typed interfaces aligned to mycli."""
+
 # mypy: ignore-errors
 
-from __future__ import print_function
-from __future__ import unicode_literals
+from __future__ import annotations
+
 import logging
 from re import compile, escape
 from collections import Counter
+from typing import Any, Collection, Generator, Iterable, List, Optional, Set, Tuple, Literal
 
-from prompt_toolkit.completion import Completer, Completion
+from prompt_toolkit.completion import CompleteEvent, Completer, Completion
+from prompt_toolkit.completion.base import Document
 
 from .packages.completion_engine import suggest_type
 from .packages.parseutils import last_word
@@ -18,7 +22,7 @@
 
 
 class SQLCompleter(Completer):
-    keywords = [
+    keywords: List[str] = [
         "ABORT",
         "ACTION",
         "ADD",
@@ -184,7 +188,7 @@ class SQLCompleter(Completer):
         "WITHOUT",
     ]
 
-    functions = [
+    functions: List[str] = [
         "ABS",
         "AVG",
         "CHANGES",
@@ -255,49 +259,49 @@ class SQLCompleter(Completer):
         "TRIM",
     ]
 
-    def __init__(self, supported_formats=(), keyword_casing="auto"):
+    def __init__(self, supported_formats: Iterable[str] = (), keyword_casing: Literal["upper", "lower", "auto"] = "auto"):
         super(self.__class__, self).__init__()
-        self.reserved_words = set()
+        self.reserved_words: Set[str] = set()
         for x in self.keywords:
             self.reserved_words.update(x.split())
         self.name_pattern = compile(r"^[_a-zA-Z][_a-zA-Z0-9\$]*$")
 
-        self.special_commands = []
-        self.table_formats = supported_formats
+        self.special_commands: List[str] = []
+        self.table_formats: List[str] = list(supported_formats)
         if keyword_casing not in ("upper", "lower", "auto"):
             keyword_casing = "auto"
-        self.keyword_casing = keyword_casing
+        self.keyword_casing: Literal["upper", "lower", "auto"] = keyword_casing
         self.reset_completions()
 
-    def escape_name(self, name):
+    def escape_name(self, name: str) -> str:
         if name and ((not self.name_pattern.match(name)) or (name.upper() in self.reserved_words) or (name.upper() in self.functions)):
             name = "`%s`" % name
 
         return name
 
-    def unescape_name(self, name):
+    def unescape_name(self, name: str) -> str:
         """Unquote a string."""
         if name and name[0] == '"' and name[-1] == '"':
             name = name[1:-1]
 
         return name
 
-    def escaped_names(self, names):
+    def escaped_names(self, names: Iterable[str]) -> List[str]:
         return [self.escape_name(name) for name in names]
 
-    def extend_special_commands(self, special_commands):
+    def extend_special_commands(self, special_commands: Iterable[str]) -> None:
         # Special commands are not part of all_completions since they can only
         # be at the beginning of a line.
         self.special_commands.extend(special_commands)
 
-    def extend_database_names(self, databases):
+    def extend_database_names(self, databases: Iterable[str]) -> None:
         self.databases.extend(databases)
 
-    def extend_keywords(self, additional_keywords):
+    def extend_keywords(self, additional_keywords: Iterable[str]) -> None:
         self.keywords.extend(additional_keywords)
         self.all_completions.update(additional_keywords)
 
-    def extend_schemata(self, schema):
+    def extend_schemata(self, schema: Optional[str]) -> None:
         if schema is None:
             return
         metadata = self.dbmetadata["tables"]
@@ -308,7 +312,7 @@ def extend_schemata(self, schema):
             metadata[schema] = {}
         self.all_completions.update(schema)
 
-    def extend_relations(self, data, kind):
+    def extend_relations(self, data: Iterable[Iterable[str]], kind: str) -> None:
         """Extend metadata for tables or views
 
         :param data: list of (rel_name, ) tuples
@@ -340,7 +344,7 @@ def extend_relations(self, data, kind):
                 )
             self.all_completions.add(relname[0])
 
-    def extend_columns(self, column_data, kind):
+    def extend_columns(self, column_data: Iterable[Iterable[str]], kind: str) -> None:
         """Extend column metadata
 
         :param column_data: list of (rel_name, column_name) tuples
@@ -362,7 +366,7 @@ def extend_columns(self, column_data, kind):
             metadata[self.dbname][relname].append(column)
             self.all_completions.add(column)
 
-    def extend_functions(self, func_data):
+    def extend_functions(self, func_data: Iterable[Iterable[str]]) -> None:
         # 'func_data' is a generator object. It can throw an exception while
         # being consumed. This could happen if the user has launched the app
         # without specifying a database name. This exception must be handled to
@@ -381,24 +385,24 @@ def extend_functions(self, func_data):
             metadata[self.dbname][func[0]] = None
             self.all_completions.add(func[0])
 
-    def set_dbname(self, dbname):
+    def set_dbname(self, dbname: str) -> None:
         self.dbname = dbname
 
-    def reset_completions(self):
-        self.databases = []
+    def reset_completions(self) -> None:
+        self.databases: List[str] = []
         self.dbname = ""
-        self.dbmetadata = {"tables": {}, "views": {}, "functions": {}}
-        self.all_completions = set(self.keywords + self.functions)
+        self.dbmetadata: dict[str, Any] = {"tables": {}, "views": {}, "functions": {}}
+        self.all_completions: Set[str] = set(self.keywords + self.functions)
 
     @staticmethod
     def find_matches(
-        text,
-        collection,
-        start_only=False,
-        fuzzy=True,
-        casing=None,
-        punctuations="most_punctuations",
-    ):
+        text: str,
+        collection: Collection[str],
+        start_only: bool = False,
+        fuzzy: bool = True,
+        casing: Optional[str] = None,
+        punctuations: str = "most_punctuations",
+    ) -> Generator[Completion, None, None]:
         """Find completion matches for the given text.
 
         Given the user's input text and a collection of available
@@ -434,16 +438,20 @@ def find_matches(
         if casing == "auto":
             casing = "lower" if last and last[-1].islower() else "upper"
 
-        def apply_case(kw):
+        def apply_case(kw: str) -> str:
             if casing == "upper":
                 return kw.upper()
             return kw.lower()
 
         return (Completion(z if casing is None else apply_case(z), -len(text)) for x, y, z in sorted(completions))
 
-    def get_completions(self, document, complete_event):
+    def get_completions(
+        self,
+        document: Document,
+        complete_event: Optional[CompleteEvent],
+    ) -> Iterable[Completion]:
         word_before_cursor = document.get_word_before_cursor(WORD=True)
-        completions = []
+        completions: List[Completion] = []
         suggestions = suggest_type(document.text, document.text_before_cursor)
 
         for suggestion in suggestions:
@@ -551,7 +559,7 @@ def get_completions(self, document, complete_event):
 
         return completions
 
-    def find_files(self, word):
+    def find_files(self, word: str) -> Generator[Completion, None, None]:
         """Yield matching directory or file names.
 
         :param word:
@@ -565,7 +573,7 @@ def find_files(self, word):
             if suggestion:
                 yield Completion(suggestion, position)
 
-    def populate_scoped_cols(self, scoped_tbls):
+    def populate_scoped_cols(self, scoped_tbls: List[Tuple[Optional[str], str, Optional[str]]]) -> List[str]:
         """Find all columns in a set of scoped_tables
         :param scoped_tbls: list of (schema, table, alias) tuples
         :return: list of column names
@@ -603,7 +611,7 @@ def populate_scoped_cols(self, scoped_tbls):
 
         return columns
 
-    def populate_schema_objects(self, schema, obj_type):
+    def populate_schema_objects(self, schema: Optional[str], obj_type: str) -> List[str]:
         """Returns list of tables or functions for a (optional) schema"""
         metadata = self.dbmetadata[obj_type]
         schema = schema or self.dbname