Refactor code for improved type handling and error management across multiple modules

Author: tbotnz

netpalm/backend/core/confload/confload.py

@@ -169,8 +169,8 @@ def ensure_non_empty_bootstrap(cls, v: str) -> str:
     @classmethod
     def parse_webhook_headers(cls, v: Any) -> dict[str, str]:
         if isinstance(v, str):
-            return json.loads(v)
-        return v
+            return json.loads(v)  # type: ignore[no-any-return]
+        return v  # type: ignore[no-any-return]
 
     def model_post_init(self, __context: Any) -> None:
         self.whitelist = DeviceWhitelist(self.device_whitelist)

netpalm/backend/core/manager/netpalm_manager.py

@@ -18,11 +18,10 @@
 
 from netpalm.backend.core.cache.store import CacheStore
 from netpalm.backend.core.confload.confload import NetpalmSettings, get_settings
+from pydantic import BaseModel
+
 from netpalm.backend.core.models.models import (
-    GetConfig,
     QueueStrategy,
-    Script,
-    SetConfig,
 )
 from netpalm.backend.core.queue.broker import QueueBroker
 from netpalm.backend.core.queue.broker import TaskResponse as BrokerTaskResponse
@@ -51,7 +50,7 @@ def __init__(
 
     # ── task operations ───────────────────────────────────────────────────────
 
-    async def get_config(self, request: GetConfig) -> dict[str, Any]:
+    async def get_config(self, request: BaseModel) -> dict[str, Any]:
         """
         Enqueue a getconfig job.
         If cache is enabled and a result exists, return it without enqueuing.
@@ -66,7 +65,7 @@ async def get_config(self, request: GetConfig) -> dict[str, Any]:
             cached = self._cache.get(cache_key)
             if cached is not None:
                 log.debug(f"NetpalmManager.get_config: cache hit for {cache_key}")
-                return cached
+                return cached  # type: ignore[no-any-return]
 
         strategy = req_data.get("queue_strategy", QueueStrategy.fifo)
         pinned_host = host if strategy == QueueStrategy.pinned else None
@@ -79,7 +78,7 @@ async def get_config(self, request: GetConfig) -> dict[str, Any]:
         )
         return _task_to_response(task)
 
-    async def set_config(self, request: SetConfig) -> dict[str, Any]:
+    async def set_config(self, request: BaseModel) -> dict[str, Any]:
         req_data = request.model_dump(exclude_none=True)
         conn = req_data.get("connection_args", {})
         host = conn.get("host", "")
@@ -98,7 +97,7 @@ async def set_config(self, request: SetConfig) -> dict[str, Any]:
         )
         return _task_to_response(task)
 
-    async def execute_script(self, request: Script) -> dict[str, Any]:
+    async def execute_script(self, request: BaseModel) -> dict[str, Any]:
         req_data = request.model_dump(exclude_none=True)
         strategy = req_data.get("queue_strategy", QueueStrategy.fifo)
         pinned_host = req_data.get("script") if strategy == QueueStrategy.pinned else None

netpalm/backend/core/operations/getconfig.py

@@ -30,8 +30,8 @@ def execute(
         driver_obj = driver_cls(**kwargs)
         sesh = driver_obj.connect()
 
-        commandlst = [command] if isinstance(command, str) else command
-        result = driver_obj.sendcommand(sesh, commandlst) if commandlst else driver_obj.sendcommand(sesh)
+        commandlst = [command] if isinstance(command, str) else (command or [])
+        result = driver_obj.sendcommand(sesh, commandlst)
 
         if post_checks:
             run_checks(driver_obj, sesh, post_checks, "PostCheck")

netpalm/backend/core/operations/script.py

@@ -66,4 +66,4 @@ def execute(
         if webhook:
             exec_webhook_func(jobdata={"task_result": result}, webhook_payload=webhook)
 
-        return result
+        return result  # type: ignore[no-any-return]

netpalm/backend/core/operations/setconfig.py

@@ -29,7 +29,7 @@ def execute(
         settings: NetpalmSettings,
     ) -> dict[str, Any]:
         library = kwargs.get("library", "")
-        config = kwargs.get("config")
+        config: str | list[str] = kwargs.get("config", "")
         j2conf = kwargs.get("j2config")
         webhook = kwargs.get("webhook")
         pre_checks = kwargs.get("pre_checks")

netpalm/backend/core/utilities/rediz_meta.py

@@ -8,16 +8,17 @@
 from __future__ import annotations
 
 import logging
+from typing import NoReturn
 
 log = logging.getLogger(__name__)
 
 
-def write_meta_error(exception: Exception) -> None:
+def write_meta_error(exception: Exception) -> NoReturn:
     """Re-raise the exception so the executor can handle it."""
     log.exception("write_meta_error: driver error")
     raise exception
 
 
-def write_meta_error_string(data: str) -> None:
+def write_meta_error_string(data: str) -> NoReturn:
     """Raise a plain exception with the given message."""
     raise Exception(f"failed: {data}")

netpalm/backend/core/utilities/universal_template_mgr/unvrsl.py

@@ -28,10 +28,14 @@ def add_template(self, payload: dict[str, str]):
             with open(template_path, "w") as file:
                 file.write(raw_base)
             reload_extensibles_func()
-            resultdata = ResponseBasic(status="success", data={"task_result": {"added": payload["name"]}}).model_dump()
+            resultdata = ResponseBasic(
+                status=TaskResponseEnum.success, data={"task_result": {"added": payload["name"]}}
+            ).model_dump()
             return resultdata
         except Exception as e:
-            error = ResponseBasic(status="error", data={"task_result": {"error": str(e)}}).model_dump()
+            error = ResponseBasic(
+                status=TaskResponseEnum.error, data={"task_result": {"error": str(e)}}
+            ).model_dump()
             return error
 
     def remove_template(self, payload: dict[str, str]):
@@ -43,12 +47,14 @@ def remove_template(self, payload: dict[str, str]):
             )
             os.remove(template_path)
             resultdata = ResponseBasic(
-                status="success", data={"task_result": {"removed": payload["name"]}}
+                status=TaskResponseEnum.success, data={"task_result": {"removed": payload["name"]}}
             ).model_dump()
             reload_extensibles_func()
             return resultdata
         except Exception as e:
-            error = ResponseBasic(status="error", data={"task_result": {"error": str(e)}}).model_dump()
+            error = ResponseBasic(
+                status=TaskResponseEnum.error, data={"task_result": {"error": str(e)}}
+            ).model_dump()
             return error
 
     def get_template(self, payload: dict[str, str]):
@@ -63,9 +69,11 @@ def get_template(self, payload: dict[str, str]):
                 result = file.read()
             raw_base = base64.b64encode(result.encode("utf-8"))
             resultdata = ResponseBasic(
-                status="success", data={"task_result": {"base64_payload": raw_base}}
+                status=TaskResponseEnum.success, data={"task_result": {"base64_payload": raw_base}}
             ).model_dump()
             return resultdata
         except Exception as e:
-            error = ResponseBasic(status="error", data={"task_result": {"error": str(e)}}).model_dump()
+            error = ResponseBasic(
+                status=TaskResponseEnum.error, data={"task_result": {"error": str(e)}}
+            ).model_dump()
             return error

netpalm/backend/core/utilities/webhook/webhook.py

@@ -15,7 +15,7 @@ def __init__(self, whook_payload: dict):
         self.webhook_args = whook_payload.get("args", False)
         if not self.webhook_raw_name:
             self.webhook_name = config.default_webhook_name
-        self.webhook_name = self.webhook_dir_path.replace("/", ".") + self.webhook_raw_name
+        self.webhook_name = self.webhook_dir_path.replace("/", ".") + (self.webhook_raw_name or "")
         self.webhook_j2_name = whook_payload.get("j2template")
 
     def webhook_exec(self, job_data: dict):

netpalm/backend/plugins/drivers/napalm/napalm_drvr.py

@@ -78,7 +78,6 @@ def config(
 
     def logout(self, session: Any) -> None:
         try:
-            response = session.close()
-            return response
+            session.close()
         except Exception as e:
             write_meta_error(e)

netpalm/backend/plugins/drivers/ncclient/ncclient_drvr.py

@@ -132,7 +132,6 @@ def config(
 
     def logout(self, session: Any) -> None:
         try:
-            response = session.close_session()
-            return response
+            session.close_session()
         except Exception as e:
             write_meta_error(e)