Logging and docstrings

clab2drawio.py

@@ -1,60 +1,81 @@
 from cli.parser_clab2drawio import parse_arguments
 from core.diagram.custom_drawio import CustomDrawioDiagram
-from core.models.node import Node
-from core.models.link import Link
 from core.grafana.grafana_manager import GrafanaDashboard
 from core.utils.yaml_processor import YAMLProcessor
-from core.data.topology_loader import TopologyLoader
+from core.data.topology_loader import TopologyLoader, TopologyLoaderError
 from core.data.node_link_builder import NodeLinkBuilder
 from core.data.graph_level_manager import GraphLevelManager
 from core.layout.vertical_layout import VerticalLayout
 from core.layout.horizontal_layout import HorizontalLayout
-from core.config.theme_manager import ThemeManager
+from core.config.theme_manager import ThemeManager, ThemeManagerError
 from core.interactivity.interactive_manager import InteractiveManager
 from core.diagram.diagram_builder import DiagramBuilder
-from collections import defaultdict
+from core.logging_config import configure_logging
 from prompt_toolkit.shortcuts import checkboxlist_dialog, yes_no_dialog
 import os
+import sys
+import logging
 
+logger = logging.getLogger(__name__)
 
 def main(
-    input_file,
-    output_file,
-    grafana,
-    theme,
-    include_unlinked_nodes=False,
-    no_links=False,
-    layout="vertical",
-    verbose=False,
-    interactive=False,
-):
+    input_file: str,
+    output_file: str,
+    grafana: bool,
+    theme: str,
+    include_unlinked_nodes: bool=False,
+    no_links: bool=False,
+    layout: str="vertical",
+    verbose: bool=False,
+    interactive: bool=False,
+) -> None:
+    """
+    Main function to generate a topology diagram from a containerlab YAML or draw.io XML file.
+
+    :param input_file: Path to the containerlab YAML file.
+    :param output_file: Output file path for the generated diagram.
+    :param grafana: Whether to generate Grafana dashboard artifacts.
+    :param theme: Theme name or path to a custom theme file.
+    :param include_unlinked_nodes: Include nodes without any links in the topology diagram.
+    :param no_links: Do not draw links between nodes.
+    :param layout: Layout direction ("vertical" or "horizontal").
+    :param verbose: Enable verbose output.
+    :param interactive: Run in interactive mode to define graph-levels and icons.
+    """    
+    logger.debug("Starting clab2drawio main function.")
     loader = TopologyLoader()
-    containerlab_data = loader.load(input_file)
+    try:
+        containerlab_data = loader.load(input_file)
+    except TopologyLoaderError as e:
+        logger.error("Failed to load topology. Exiting.")
+        sys.exit(1)
 
     try:
         if os.path.isabs(theme):
             theme_path = theme
         else:
             theme_path = os.path.join(script_dir, "styles", f"{theme}.yaml")
 
-        # Check if the theme file exists
         if not os.path.exists(theme_path):
             raise FileNotFoundError(
                 f"The specified theme file '{theme_path}' does not exist."
             )
-
     except FileNotFoundError as e:
-        error_message = str(e)
-        print(error_message)
-        exit()
+        logger.error(str(e))
+        sys.exit(1)
     except Exception as e:
-        error_message = f"An error occurred while loading the theme: {e}"
-        print(error_message)
-        exit()
+        logger.error(f"An error occurred while loading the theme: {e}")
+        sys.exit(1)
 
     # Use ThemeManager to load styles
+    logger.debug("Loading theme...")
     theme_manager = ThemeManager(theme_path)
-    styles = theme_manager.load_theme()
+    try:
+        styles = theme_manager.load_theme()
+    except ThemeManagerError as e:
+        logger.error("Failed to load theme. Exiting.")
+        sys.exit(1)
+    logger.debug("Theme loaded successfully, building diagram...")
 
     diagram = CustomDrawioDiagram()
     diagram.layout = layout
@@ -65,6 +86,7 @@ def main(
     lab_name = containerlab_data.get("name", "")
 
     # Use NodeLinkBuilder to build nodes and links
+    logger.debug("Building nodes and links...")
     builder = NodeLinkBuilder(containerlab_data, styles, prefix, lab_name)
     nodes, links = builder.build_nodes_and_links()    
 
@@ -78,6 +100,7 @@ def main(
         diagram.nodes = nodes
 
     if interactive:
+        logger.debug("Entering interactive mode...")
         processor = YAMLProcessor()
         interactor = InteractiveManager()
         interactor.run_interactive_mode(
@@ -90,6 +113,7 @@ def main(
             lab_name,
         )
 
+    logger.debug("Assigning graph levels...")
     graph_manager = GraphLevelManager()
     graph_manager.assign_graphlevels(diagram, verbose=False)
 
@@ -99,6 +123,7 @@ def main(
     else:
         layout_manager = HorizontalLayout()
 
+    logger.debug(f"Applying {layout} layout...")
     layout_manager.apply(diagram, verbose=verbose)
 
     # Calculate the diagram size based on the positions of the nodes
@@ -128,17 +153,20 @@ def main(
     if styles["pageh"] == "auto":
         styles["pageh"] = max_size_y
 
+    logger.debug("Updating diagram style...")
     diagram.update_style(styles)
 
     diagram.add_diagram("Network Topology")
 
     diagram_builder = DiagramBuilder()
+    logger.debug("Adding nodes to diagram...")
     diagram_builder.add_nodes(diagram, diagram.nodes, styles)
 
     if grafana:
         styles["ports"] = True
 
     if styles["ports"]:
+        logger.debug("Adding ports and generating Grafana dashboard...")
         diagram_builder.add_ports(diagram, styles)
         if not output_file:
             grafana_output_file = os.path.splitext(input_file)[0] + ".grafana.json"
@@ -159,6 +187,7 @@ def main(
             f.write(grafana_json)
         print("Saved Grafana dashboard JSON to:", grafana_output_file)
     else:
+        logger.debug("Adding links to diagram...")
         diagram_builder.add_links(diagram, styles)
 
     if not output_file:
@@ -168,6 +197,7 @@ def main(
     output_filename = os.path.basename(output_file)
     os.makedirs(output_folder, exist_ok=True)
 
+    logger.debug(f"Dumping diagram to file: {output_file}")
     diagram.dump_file(filename=output_filename, folder=output_folder)
 
     print("Saved file to:", output_file)
@@ -177,6 +207,10 @@ def main(
 
     script_dir = os.path.dirname(__file__)
 
+    # Configure logging at startup
+    log_level = logging.DEBUG if args.verbose else logging.INFO
+    configure_logging(level=log_level)
+
     main(
         args.input,
         args.output,

cli/parser_clab2drawio.py

@@ -1,6 +1,11 @@
 import argparse
 
 def parse_arguments():
+    """
+    Parse command-line arguments for clab2drawio tool.
+
+    :return: argparse.Namespace object with parsed arguments.
+    """
     parser = argparse.ArgumentParser(
         description="Generate a topology diagram from a containerlab YAML or draw.io XML file."
     )

cli/parser_drawio2clab.py

@@ -1,6 +1,11 @@
 import argparse
 
 def parse_arguments():
+    """
+    Parse command-line arguments for drawio2clab tool.
+
+    :return: argparse.Namespace object with parsed arguments.
+    """
     parser = argparse.ArgumentParser(
         description="Convert a .drawio file to a Containerlab YAML file."
     )

core/config/theme_manager.py

@@ -1,23 +1,41 @@
 import yaml
 import sys
+import logging
+
+logger = logging.getLogger(__name__)
+
+class ThemeManagerError(Exception):
+    """Raised when loading the theme fails."""
 
 class ThemeManager:
-    def __init__(self, config_path):
+    """
+    Loads and merges theme configuration from a YAML file.
+    """
+    def __init__(self, config_path: str):
+        """
+        :param config_path: Path to the theme configuration file.
+        """
         self.config_path = config_path
 
-    def load_theme(self):
+    def load_theme(self) -> dict:
+        """
+        Load the theme configuration and return a dictionary of styles.
+
+        :return: Dictionary containing styles and configuration parameters.
+        """
+        logger.debug(f"Loading theme from: {self.config_path}")
         try:
             with open(self.config_path, "r") as file:
                 config = yaml.safe_load(file)
         except FileNotFoundError:
             error_message = (
                 f"Error: The specified config file '{self.config_path}' does not exist."
             )
-            print(error_message)
+            logger.error(error_message)
             sys.exit(1)
         except Exception as e:
             error_message = f"An error occurred while loading the config: {e}"
-            print(error_message)
+            logger.error(error_message)
             sys.exit(1)
 
         base_style_dict = {
@@ -42,7 +60,6 @@ def load_theme(self):
             "custom_styles": {},
         }
 
-        # Merge base style with custom styles
         for key, custom_style in config.get("custom_styles", {}).items():
             custom_style_dict = {
                 item.split("=")[0]: item.split("=")[1]
@@ -53,7 +70,6 @@ def load_theme(self):
             merged_style = ";".join(f"{k}={v}" for k, v in merged_style_dict.items())
             styles["custom_styles"][key] = merged_style
 
-        # Read all other configuration values
         for key, value in config.items():
             if key not in styles:
                 styles[key] = value

core/data/graph_level_manager.py

@@ -1,9 +1,17 @@
-class GraphLevelManager:
-    def __init__(self):
-        pass
+import logging
+
+logger = logging.getLogger(__name__)
 
-    def update_links(self, links):
-        # Same logic as before
+class GraphLevelManager:
+    """
+    Manages graph-level assignments and adjustments for nodes in the diagram.
+    """
+
+    def update_links(self, links: list) -> None:
+        """
+        Update the direction and level difference of links after changes in node graph levels.
+        """
+        logger.debug("Updating link directions and level differences...")
         for link in links:
             source_level = link.source.graph_level
             target_level = link.target.graph_level
@@ -15,13 +23,17 @@ def update_links(self, links):
             else:
                 link.direction = "lateral"
 
-    def adjust_node_levels(self, diagram):
+    def adjust_node_levels(self, diagram) -> None:
+        """
+        Adjust node levels to reduce inconsistencies and improve layout clarity.
+        """
+        logger.debug("Adjusting node levels for better layout...")
         used_levels = diagram.get_used_levels()
         max_level = diagram.get_max_level()
         min_level = diagram.get_min_level()
 
         if len(used_levels) <= 1:
-            return  # Only one level present, no adjustment needed
+            return
 
         current_level = min_level
         while current_level < max_level + 1:
@@ -83,15 +95,23 @@ def adjust_node_levels(self, diagram):
                             max_level = diagram.get_max_level()
                             break
 
-    def assign_graphlevels(self, diagram, verbose=False):
+    def assign_graphlevels(self, diagram, verbose=False) -> list:
+        """
+        Assign initial graph levels to nodes and adjust them as necessary.
+
+        :param diagram: CustomDrawioDiagram instance.
+        :param verbose: Whether to print debugging info.
+        :return: Sorted list of nodes by level and name.
+        """
+        logger.debug("Assigning graph levels to nodes...")
         nodes = diagram.get_nodes()
 
-        # Check if all nodes already have a graphlevel != -1
         if all(node.graph_level != -1 for node in nodes.values()):
             already_set = True
         else:
             already_set = False
-            print(
+            if verbose:
+                print(
                 "Not all graph levels set in the .clab file. Assigning graph levels based on downstream links. Expect experimental output. Please consider assigning graph levels to your .clab file, or use it with -I for interactive mode. Find more information here: https://github.com/srl-labs/clab-io-draw/blob/grafana_style/docs/clab2drawio.md#influencing-node-placement"
             )
 
@@ -127,4 +147,5 @@ def set_graphlevel(node, current_graphlevel, visited=None):
         sorted_nodes = sorted(
             nodes.values(), key=lambda node: (node.graph_level, node.name)
         )
+        logger.debug("Graph levels assigned.")
         return sorted_nodes