fix: resolve ty optional import errors

Author: amjith

CHANGELOG.md

@@ -8,6 +8,10 @@
 
 - Avoid completion refresh crashes when no database is connected.
 
+### Internal
+
+- Clean up ty type-checking for optional sqlean/llm imports.
+
 ## 1.18.0 
 
 ### Internal

litecli/__init__.py

@@ -1,4 +1,3 @@
-# type: ignore
 from __future__ import annotations
 
 import importlib.metadata

litecli/main.py

@@ -11,11 +11,6 @@
 from collections import namedtuple
 from datetime import datetime
 from io import open
-
-try:
-    from sqlean import OperationalError, sqlite_version
-except ImportError:
-    from sqlite3 import OperationalError, sqlite_version
 from time import time
 from typing import Any, Generator, Iterable, cast
 
@@ -51,6 +46,15 @@
 from .sqlcompleter import SQLCompleter
 from .sqlexecute import SQLExecute
 
+try:
+    import sqlean as _sqlite3
+except ImportError:
+    import sqlite3 as _sqlite3
+
+_sqlite3 = cast(Any, _sqlite3)
+OperationalError = _sqlite3.OperationalError
+sqlite_version = _sqlite3.sqlite_version
+
 # Query tuples are used for maintaining history
 Query = namedtuple("Query", ["query", "successful", "mutating"])
 

litecli/packages/special/llm.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextlib
+import importlib
 import io
 import logging
 import os
@@ -15,32 +16,40 @@
 import click
 
 try:
-    import llm
-
-    LLM_IMPORTED = True
+    import llm as llm_module
 except ImportError:
-    llm = None
-    LLM_IMPORTED = False
+    llm_module = None
 
 try:
-    from llm.cli import cli
-
-    LLM_CLI_IMPORTED = True
+    llm_cli_module = importlib.import_module("llm.cli")
 except ImportError:
-    cli = None
-    LLM_CLI_IMPORTED = False
+    llm_cli_module = None
 
 from . import export
 from .main import Verbosity, parse_special_command
 from .types import DBCursor
 
+LLM_IMPORTED = llm_module is not None
+
+cli: click.Command | None
+if llm_cli_module is not None:
+    llm_cli = getattr(llm_cli_module, "cli", None)
+    cli = llm_cli if isinstance(llm_cli, click.Command) else None
+else:
+    cli = None
+
+LLM_CLI_IMPORTED = cli is not None
+
 log = logging.getLogger(__name__)
 
 LLM_TEMPLATE_NAME = "litecli-llm-template"
-LLM_CLI_COMMANDS: list[str] = list(cli.commands.keys()) if LLM_CLI_IMPORTED else []
+LLM_CLI_COMMANDS: list[str] = list(cli.commands.keys()) if isinstance(cli, click.Group) else []
 # Mapping of model_id to None used for completion tree leaves.
-# the file name is llm.py and module name is llm, hence ty is complaining that get_models is missing.
-MODELS: dict[str, None] = {x.model_id: None for x in llm.get_models()} if LLM_IMPORTED else {}  # type: ignore[attr-defined]
+if llm_module is not None:
+    get_models = getattr(llm_module, "get_models", None)
+    MODELS: dict[str, None] = {x.model_id: None for x in get_models()} if callable(get_models) else {}
+else:
+    MODELS = {}
 
 
 def run_external_cmd(
@@ -124,7 +133,7 @@ def build_command_tree(cmd: click.Command) -> dict[str, Any] | None:
 
 
 # Generate the tree
-COMMAND_TREE: dict[str, Any] | None = build_command_tree(cli) if LLM_CLI_IMPORTED else {}
+COMMAND_TREE: dict[str, Any] | None = build_command_tree(cli) if cli is not None else {}
 
 
 def get_completions(tokens: list[str], tree: dict[str, Any] | None = COMMAND_TREE) -> list[str]:

litecli/sqlexecute.py

@@ -1,25 +1,26 @@
 from __future__ import annotations
 
 import logging
+import os.path
 from contextlib import closing
-from typing import Any, Generator, Iterable
+from typing import Any, Generator, Iterable, cast
+from urllib.parse import urlparse
+
+import sqlparse
 
 try:
-    import sqlean as sqlite3
-    from sqlean import OperationalError
+    import sqlean as _sqlite3
 
-    sqlite3.extensions.enable_all()
+    _sqlite3.extensions.enable_all()
 except ImportError:
-    import sqlite3
-    from sqlite3 import OperationalError
-import os.path
-from urllib.parse import urlparse
-
-import sqlparse
+    import sqlite3 as _sqlite3
 
 from litecli.packages import special
 from litecli.packages.special.utils import check_if_sqlitedotcommand
 
+sqlite3 = cast(Any, _sqlite3)
+OperationalError = sqlite3.OperationalError
+
 _logger = logging.getLogger(__name__)
 
 # FIELD_TYPES = decoders.copy()

tests/test_main.py

@@ -243,7 +243,7 @@ def stub_terminal_size():
     shutil.get_terminal_size = stub_terminal_size  # type: ignore[assignment]
     lc = LiteCli()
     assert isinstance(lc.get_reserved_space(), int)
-    shutil.get_terminal_size = old_func  # type: ignore[assignment]
+    shutil.get_terminal_size = old_func
 
 
 @dbtest

tests/test_sqlexecute.py

@@ -1,15 +1,20 @@
 # coding=UTF-8
 
 import os
+from typing import Any, cast
 
 import pytest
 
 from .utils import assert_result_equal, dbtest, is_expanded_output, run, set_expanded_output
 
 try:
-    from sqlean import OperationalError, ProgrammingError
+    import sqlean as _sqlite3
 except ImportError:
-    from sqlite3 import OperationalError, ProgrammingError
+    import sqlite3 as _sqlite3
+
+_sqlite3 = cast(Any, _sqlite3)
+OperationalError = _sqlite3.OperationalError
+ProgrammingError = _sqlite3.ProgrammingError
 
 
 @dbtest