Hacky changes (so far) for unified shell + apps mypy

pyproject.toml

@@ -75,8 +75,12 @@ source = ["drunc"]
 [tool.ruff.lint]
 ignore = [
   "E501", # Don't enforce line lengths within a linting context
+  "D201", 
+  "D202", 
+  "D212"
 ]
 select = [
+  "D",
   "E", # pycodestyle errors
   "F", # check for errors using PyFlakes
   "I", # best practices for import calls
@@ -113,3 +117,18 @@ module = ["google.rpc.*"]
 [[tool.mypy.overrides]]
 ignore_missing_imports = true
 module = ["conffwk"]
+
+[[tool.mypy.overrides]]
+ignore_missing_imports = true
+module = [
+  "click_shell",
+  "confmodel_dal",
+  "flask_restful",
+  "grpc_status",
+  "kafkaopmon.*",
+  "opmonlib.*",
+  "paramiko",
+  "pytz",
+  "sh",
+  "socks",
+]

src/drunc/apps/check_np0x_cluster.py

@@ -1,13 +1,37 @@
 import os
 import subprocess
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import NotRequired, TypedDict
 
 import paramiko
 from rich import box
 from rich.console import Console
 from rich.live import Live
 from rich.table import Table
 
+
+class HostData(TypedDict):
+    """Data dictionary tracking the status of a single cluster host."""
+
+    alias: str
+    is_pingable: bool | None
+    ssh_key: str
+    key_color: str
+    cpu_count: str
+    total_cores: str
+    cpu_usage: str
+    cpu_model: str
+    model_color: str
+    ram_total: str
+    ram_usage: str
+    nfs_ok: bool
+    uptime: str
+    details: NotRequired[str]
+    ssh_key_status: NotRequired[str]
+    ssh_key_color: NotRequired[str]
+    status: NotRequired[str]
+
+
 # List of NP0x cluster hosts
 NP0X_CLUSTER_HOSTS = sorted(
     [
@@ -48,7 +72,7 @@
 STYLE_OFFLINE = "[bold white on #444444] OFFLINE [/]"
 
 # Base template for host tracking
-DEFAULT_HOST_DATA = {
+DEFAULT_HOST_DATA: HostData = {
     "alias": "",
     "is_pingable": None,  # None: Waiting, True: Up, False: Down
     "ssh_key": "Pending",
@@ -65,18 +89,18 @@
 }
 
 
-class TrackingAutoAddPolicy(paramiko.MissingHostKeyPolicy):
+class TrackingAutoAddPolicy(paramiko.MissingHostKeyPolicy):  # type: ignore[misc, no-any-unimported]
     """
     Custom policy to track missing host keys and update the result dict accordingly.
     """
 
-    def __init__(self, res_dict):
+    def __init__(self, res_dict: HostData) -> None:
         """
         Initialize with a reference to the result dictionary to update key status.
         """
         self.res_dict = res_dict
 
-    def missing_host_key(self, client, hostname, key):
+    def missing_host_key(self, client: object, hostname: str, key: object) -> None:
         """
         When a host key is missing, update the result dictionary to indicate that the
         key is being added.
@@ -86,7 +110,7 @@ def missing_host_key(self, client, hostname, key):
         self.res_dict["key_color"] = "bold yellow"
 
 
-def load_ssh_config() -> paramiko.SSHConfig:
+def load_ssh_config() -> paramiko.SSHConfig:  # type: ignore[no-any-unimported]
     """
     Load the user's SSH configuration from ~/.ssh/config using Paramiko's SSHConfig
     class.
@@ -146,7 +170,7 @@ def ping_host(hostname: str) -> bool:
     )
 
 
-def get_host_metrics(host_alias: str, ssh_config: paramiko.SSHConfig) -> dict:
+def get_host_metrics(host_alias: str, ssh_config: paramiko.SSHConfig) -> HostData:  # type: ignore[no-any-unimported]
     """
     For a given host alias, perform the following checks:
          - Ping the host to determine if it's online.
@@ -326,7 +350,7 @@ def get_host_metrics(host_alias: str, ssh_config: paramiko.SSHConfig) -> dict:
     return result
 
 
-def generate_table(results_map: dict[str, str]) -> Table:
+def generate_table(results_map: dict[str, HostData]) -> Table:
     """
     Generate a Rich Table object to display the status of the NP0x cluster hosts.
 
@@ -372,8 +396,8 @@ def generate_table(results_map: dict[str, str]) -> Table:
     # Iterate through the results_map and add a row to the table for each host.
     for host in NP0X_CLUSTER_HOSTS:
         # Retrieve the result dictionary for the current host
-        r = results_map.get(host)
-        hostname = r["alias"] if r else host
+        r = results_map.get(host, DEFAULT_HOST_DATA)
+        hostname = r["alias"] if r["alias"] else host
 
         # Host is Pingable
         if r["is_pingable"] is True:
@@ -415,7 +439,7 @@ def generate_table(results_map: dict[str, str]) -> Table:
     return table
 
 
-def main():
+def main() -> None:
     """
     Main function to execute the NP0x cluster status check and display results in a
     live-updating table.

src/drunc/apps/check_np0x_hw_status.py

@@ -11,6 +11,7 @@
 import sys
 from concurrent.futures import ThreadPoolExecutor, TimeoutError, as_completed
 from datetime import datetime
+from typing import TypedDict
 
 import click
 import pytz
@@ -21,8 +22,15 @@
 from rich.table import Table
 
 
+class HardwareStatus(TypedDict):
+    """Status dictionary returned by check_hardware."""
+
+    online: bool
+    fembs: list[bool]
+
+
 # --- DYNAMIC PATH SETUP ---
-def setup_wib_path():
+def setup_wib_path() -> str | None:
     """
     Sets up the path to the local copy of the dune-wib-firmware repository if it exists.
 
@@ -67,8 +75,8 @@ def setup_wib_path():
 
 if WIB_FW_SW_IFACE_PATH:
     try:
-        import wib_pb2 as wibpb
-        from wib import WIB
+        import wib_pb2 as wibpb  # type: ignore[import-not-found, no-redef]
+        from wib import WIB  # type: ignore[import-not-found, no-redef]
 
         WIB_LIB_AVAILABLE = True
     except (ImportError, ModuleNotFoundError):
@@ -156,7 +164,7 @@ def setup_wib_path():
 }
 
 
-def check_hardware(ip: str) -> dict:
+def check_hardware(ip: str) -> HardwareStatus:
     """
     Checks the hardware status of a given IP address by first pinging it to determine if
     it is online, and if it is online, attempting to query its FEMB power status using
@@ -180,7 +188,7 @@ def check_hardware(ip: str) -> dict:
             errors (e.g. timeouts, subprocess errors, gRPC errors, etc.).
     """
     # Default state is 'In Progress' (None)
-    final_status = {"online": False, "fembs": [False] * 4}
+    final_status: HardwareStatus = {"online": False, "fembs": [False] * 4}
 
     try:
         # 1. Ping
@@ -193,7 +201,7 @@ def check_hardware(ip: str) -> dict:
             final_status["online"] = True
 
             # 2. Protocol Check
-            if WIB_LIB_AVAILABLE:
+            if WIB is not None and wibpb is not None:
                 try:
                     wib_inst = WIB(ip)
                     req = wibpb.GetFEMBStatus()
@@ -214,7 +222,9 @@ def check_hardware(ip: str) -> dict:
     return final_status
 
 
-def make_wib_table(category: str, wibs: dict[str, str], results_map: dict) -> Table:
+def make_wib_table(
+    category: str, wibs: dict[str, str], results_map: dict[str, HardwareStatus | None]
+) -> Table:
     """
     Creates a Rich Table for a given category of WIBs, showing their online status and
     FEMB power status.
@@ -281,7 +291,7 @@ def make_wib_table(category: str, wibs: dict[str, str], results_map: dict) -> Ta
     return table
 
 
-def generate_display(results_map: dict) -> Table:
+def generate_display(results_map: dict[str, HardwareStatus | None]) -> Table:
     """
     Generates the overall display grid for the current results.
 
@@ -326,7 +336,7 @@ def generate_display(results_map: dict) -> Table:
         "Each WIB is pinged first, then queried for FEMB status when reachable."
     ),
 )
-def main():
+def main() -> None:
     """
     Prints the power status of the hardware defined in WIB_DATA.
 
@@ -355,7 +365,7 @@ def main():
     ]
 
     # Initialize results map as empty/None for all IPs to force dots initially
-    results = {ip: None for ip in all_ips}
+    results: dict[str, HardwareStatus | None] = {ip: None for ip in all_ips}
 
     # Define a ThreadPoolExecutor to check hardware in parallel, and a mapping of
     # futures to IPs

src/drunc/apps/controller.py

@@ -2,7 +2,7 @@
 from drunc.utils.utils import get_logger, get_root_logger
 
 
-def main():
+def main() -> None:
     try:
         controller_cli()
     except Exception as e: