add docstrings

diode-napalm-agent/diode_napalm/__init__.py

@@ -1,3 +1,3 @@
 #!/usr/bin/env python
 # Copyright 2024 NetBox Labs Inc
-"""NetBox Labs - Diode NAPAML namespace."""
+"""NetBox Labs - Diode NAPALM namespace."""

diode-napalm-agent/diode_napalm/cli/cli.py

@@ -1,10 +1,11 @@
 
 #!/usr/bin/env python
 # Copyright 2024 NetBox Labs Inc
-"""Diode NAPALM Agent CLI"""
+"""Diode NAPALM Agent CLI."""
 
 import argparse
 import asyncio
+import logging
 import sys
 from importlib.metadata import version
 from pathlib import Path
@@ -15,8 +16,24 @@
 from diode_napalm.version import version_semver
 import netboxlabs.diode.sdk.version as SdkVersion
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 
 def parse_config_file(cfile: Path):
+    """
+    Parse the configuration file and return the diode configuration.
+
+    Args:
+        cfile (Path): Path to the configuration file.
+
+    Returns:
+        cfg.diode: Parsed configuration data.
+
+    Raises:
+        SystemExit: If the configuration file cannot be opened or parsed.
+    """
     try:
         with open(cfile, "r") as f:
             cfg = parse_config(f.read())
@@ -28,21 +45,34 @@ def parse_config_file(cfile: Path):
 
 
 async def start_policy(cfg, client):
+    """
+    Start the policy for the given configuration and client.
+
+    Args:
+        cfg: Configuration data for the policy.
+        client: Client instance for data ingestion.
+    """
     for info in cfg.data:
-        print(f"Get driver '{info.driver}'")
+        logger.info(f"Get driver '{info.driver}'")
         np_driver = get_network_driver(info.driver)
-        print(f"Getting information from '{info.hostname}'")
+        logger.info(f"Getting information from '{info.hostname}'")
         with np_driver(info.hostname, info.username, info.password, info.timeout, info.optional_args) as device:
-            data = {}
-            data["driver"] = info.driver
-            data["site"] = cfg.config.netbox.get("site", None)
-            data["device"] = device.get_facts()
-            data["interface"] = device.get_interfaces()
+            data = {
+                "driver": info.driver,
+                "site": cfg.config.netbox.get("site", None),
+                "device": device.get_facts(),
+                "interface": device.get_interfaces(),
+            }
             client.ingest(data)
 
 
 async def start_agent(cfg):
-    # start diode client
+    """
+    Start the diode client and execute policies.
+
+    Args:
+        cfg: Configuration data containing policies.
+    """
     client = Client()
     client.init_client(target=cfg.config.target,
                        api_key=cfg.config.api_key, tls_verify=cfg.config.tls_verify)
@@ -53,14 +83,18 @@ async def start_agent(cfg):
             raise Exception(f"Unable to start policy {policy_name}: {e}")
 
 
-def agent_main():
-    parser = argparse.ArgumentParser(description="Diode Agent for SuzieQ")
+def main():
+    """
+    Main entry point for the Diode NAPALM Agent CLI.
+    Parses command-line arguments and starts the agent.
+    """
+    parser = argparse.ArgumentParser(description="Diode Agent for NAPALM")
     parser.add_argument(
         "-V",
         "--version",
         action="version",
-        version=f"Diode Agent version: {version_semver()}, SuzieQ version: {version('napalm')}, Diode SDK version: {SdkVersion.version_semver()}",
-        help='Display Diode Agent, SuzieQ and Diode SDK versions'
+        version=f"Diode Agent version: {version_semver()}, NAPALM version: {version('napalm')}, Diode SDK version: {SdkVersion.version_semver()}",
+        help='Display Diode Agent, NAPALM and Diode SDK versions'
     )
     parser.add_argument(
         "-c",
@@ -80,4 +114,4 @@ def agent_main():
 
 
 if __name__ == '__main__':
-    agent_main()
+    main()

diode-napalm-agent/diode_napalm/client.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python
 # Copyright 2024 NetBox Labs Inc
-"""Diode SDK Client for NAPALM"""
+"""Diode SDK Client for NAPALM."""
 
+import logging
 from typing import Optional
 from diode_napalm.version import version_semver
 from diode_napalm.translate import translate_data
@@ -10,16 +11,38 @@
 APP_NAME = "diode-napalm-agent"
 APP_VERSION = version_semver()
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 
 class Client:
+    """
+    Singleton class for managing the Diode client for NAPALM.
+
+    This class ensures only one instance of the Diode client is created and provides methods
+    to initialize the client and ingest data.
+
+    Attributes:
+        diode_client (DiodeClient): Instance of the DiodeClient.
+    """
     _instance = None
 
     def __new__(cls):
+        """
+        Create a new instance of the Client if one does not already exist.
+
+        Returns:
+            Client: The singleton instance of the Client.
+        """
         if cls._instance is None:
             cls._instance = super(Client, cls).__new__(cls)
         return cls._instance
 
     def __init__(self):
+        """
+        Initialize the Client instance with no Diode client.
+        """
         self.diode_client = None
 
     def init_client(
@@ -28,12 +51,31 @@ def init_client(
         api_key: Optional[str] = None,
         tls_verify: bool = None
     ):
+        """
+        Initialize the Diode client with the specified target, API key, and TLS verification.
+
+        Args:
+            target (str): The target endpoint for the Diode client.
+            api_key (Optional[str]): The API key for authentication (default is None).
+            tls_verify (bool): Whether to verify TLS certificates (default is None).
+        """
         self.diode_client = DiodeClient(
             target=target, app_name=APP_NAME, app_version=APP_VERSION, api_key=api_key, tls_verify=tls_verify)
 
     def ingest(self, data: dict):
+        """
+        Ingest data using the Diode client after translating it.
+
+        Args:
+            data (dict): The data to be ingested.
+
+        Raises:
+            ValueError: If the Diode client is not initialized.
+        """
         if self.diode_client is None:
             raise ValueError("diode client defined")
         ret = self.diode_client.ingest(translate_data(data))
-        print(ret)
-
+        if not len(ret.errors):
+            logger.info("successful ingestion")
+        else:
+            logger.error(ret)

diode-napalm-agent/diode_napalm/parser.py

@@ -1,17 +1,19 @@
 #!/usr/bin/env python
 # Copyright 2024 NetBox Labs Inc
-"""Parse Diode Agent Config file"""
+"""Parse Diode Agent Config file."""
 
 import yaml
 from typing import Any, Dict, List, Optional
 from pydantic import BaseModel, Field, ValidationError
 
 
 class ParseException(Exception):
+    """Custom exception for parsing errors."""
     pass
 
 
 class Napalm(BaseModel):
+    """Model for NAPALM configuration."""
     driver: str
     hostname: str
     username: str
@@ -22,30 +24,47 @@ class Napalm(BaseModel):
 
 
 class DiscoveryConfig(BaseModel):
+    """Model for discovery configuration."""
     netbox: Dict[str, str]
 
 
 class Policy(BaseModel):
+    """Model for a policy configuration."""
     config: DiscoveryConfig
     data: List[Napalm]
 
 
 class DiodeConfig(BaseModel):
+    """Model for Diode configuration."""
     target: str
     api_key: str
     tls_verify: bool
 
 
 class Diode(BaseModel):
+    """Model for Diode containing configuration and policies."""
     config: DiodeConfig
     policies: Dict[str, Policy]
 
 
 class Config(BaseModel):
+    """Top-level model for the entire configuration."""
     diode: Diode
 
 
 def parse_config(config_data: str):
+    """
+    Parse the YAML configuration data into a Config object.
+
+    Args:
+        config_data (str): The YAML configuration data as a string.
+
+    Returns:
+        Config: The parsed configuration object.
+
+    Raises:
+        ParseException: If there is an error in parsing the YAML or validating the data.
+    """
     try:
         config = Config(**yaml.safe_load(config_data))
         return config

diode-napalm-agent/diode_napalm/translate.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # Copyright 2024 NetBox Labs Inc
-"""Translate from NAPALM output format to Diode SDK entites"""
+"""Translate from NAPALM output format to Diode SDK entities."""
 
 from typing import Iterable
 from netboxlabs.diode.sdk.diode.v1.ingester_pb2 import (
@@ -15,6 +15,15 @@
 
 
 def translate_device(device_info: dict) -> Device:
+    """
+    Translate device information from NAPALM format to Diode SDK Device entity.
+
+    Args:
+        device_info (dict): Dictionary containing device information.
+
+    Returns:
+        Device: Translated Device entity.
+    """
     manufacturer = Manufacturer(name=device_info["vendor"])
     device_type = DeviceType(
         model=device_info["model"],
@@ -37,6 +46,17 @@ def translate_device(device_info: dict) -> Device:
 
 
 def translate_interface(device: Device, if_name: str, interface_info: dict) -> Interface:
+    """
+    Translate interface information from NAPALM format to Diode SDK Interface entity.
+
+    Args:
+        device (Device): The device to which the interface belongs.
+        if_name (str): The name of the interface.
+        interface_info (dict): Dictionary containing interface information.
+
+    Returns:
+        Interface: Translated Interface entity.
+    """
     interface = Interface(
         device=device,
         name=if_name,
@@ -51,17 +71,30 @@ def translate_interface(device: Device, if_name: str, interface_info: dict) -> I
 
 
 def translate_data(data: dict) -> Iterable[Entity]:
+    """
+    Translate data from NAPALM format to Diode SDK entities.
+
+    Args:
+        data (dict): Dictionary containing data to be translated.
+
+    Returns:
+        Iterable[Entity]: Iterable of translated entities.
+    """
     entities = []
 
-    if "device" in data:
-        data["device"]["driver"] = data["driver"]
-        data["device"]["site"] = data["site"]
-        device = translate_device(data["device"])
+    device_info = data.get("device")
+    if device_info:
+        device_info["driver"] = data.get("driver")
+        device_info["site"] = data.get("site")
+        device = translate_device(device_info)
         entities.append(Entity(device=device))
-        if "interface" in data:
-            for if_name in data["interface"]:
-                if if_name in data["device"]["interface_list"]:
-                    entities.append(Entity(interface=translate_interface(
-                        device, if_name, data["interface"][if_name])))
+
+        interfaces = data.get("interface", {})
+        interface_list = device_info.get("interface_list", [])
+        for if_name, interface_info in interfaces.items():
+            if if_name in interface_list:
+                interface = translate_interface(
+                    device, if_name, interface_info)
+                entities.append(Entity(interface=interface))
 
     return entities

diode-napalm-agent/pyproject.toml

@@ -1,10 +1,10 @@
 [project]
 name = "diode-napalm-agent"
 version = "0.0.1"  # Overwritten during the build process
-description = "NetBox Labs, Diode SuzieQ Agent"
+description = "NetBox Labs, Diode NAPALM Agent"
 readme = "README.md" 
 requires-python = ">=3.7"
-license = {file = "LICENSE.txt"}
+license = {file = "../LICENSE.txt"}
 authors = [
     {name = "NetBox Labs", email = "[email protected]" } # Optional
 ]
@@ -40,7 +40,7 @@ test = ["coverage", "pytest", "pytest-cov"]
 "Homepage" = "https://netboxlabs.com/"
 
 [project.scripts]  # Optional
-diode-napalm-agent = "diode_napalm.cli.cli:agent_main"
+diode-napalm-agent = "diode_napalm.cli.cli:main"
 
 [tool.setuptools]
 packages = [
@@ -55,10 +55,6 @@ build-backend = "setuptools.build_meta"
 
 [tool.ruff]
 line-length = 140
-exclude = [
-    "netboxlabs/diode/sdk/diode/*",
-    "netboxlabs/diode/sdk/validate/*",
-]
 
 [tool.ruff.format]
 quote-style = "double"