diff --git a/clab2drawio.py b/clab2drawio.py index 079e75f..ff9c94e 100644 --- a/clab2drawio.py +++ b/clab2drawio.py @@ -1,35 +1,54 @@ 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): @@ -37,24 +56,26 @@ def main( 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, diff --git a/cli/parser_clab2drawio.py b/cli/parser_clab2drawio.py index 7205739..ed01ac1 100644 --- a/cli/parser_clab2drawio.py +++ b/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." ) diff --git a/cli/parser_drawio2clab.py b/cli/parser_drawio2clab.py index 160196f..f246105 100644 --- a/cli/parser_drawio2clab.py +++ b/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." ) diff --git a/core/config/theme_manager.py b/core/config/theme_manager.py index 95fae3a..e57394e 100644 --- a/core/config/theme_manager.py +++ b/core/config/theme_manager.py @@ -1,11 +1,29 @@ 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) @@ -13,11 +31,11 @@ def load_theme(self): 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 diff --git a/core/data/graph_level_manager.py b/core/data/graph_level_manager.py index b65be7a..52a13ba 100644 --- a/core/data/graph_level_manager.py +++ b/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 diff --git a/core/data/node_link_builder.py b/core/data/node_link_builder.py index 160c280..e270e98 100644 --- a/core/data/node_link_builder.py +++ b/core/data/node_link_builder.py @@ -2,7 +2,16 @@ from core.models.link import Link class NodeLinkBuilder: - def __init__(self, containerlab_data, styles, prefix, lab_name): + """ + Builds Node and Link objects from containerlab topology data and styling information. + """ + def __init__(self, containerlab_data: dict, styles: dict, prefix: str, lab_name: str): + """ + :param containerlab_data: Parsed containerlab topology data. + :param styles: Dictionary of style parameters. + :param prefix: Prefix used in node names. + :param lab_name: Name of the lab. + """ self.containerlab_data = containerlab_data self.styles = styles self.prefix = prefix @@ -17,11 +26,23 @@ def format_node_name(self, base_name): return f"{self.prefix}-{self.lab_name}-{base_name}" def build_nodes_and_links(self): + """ + Build Node and Link objects from the provided containerlab data and return them. + + :return: A tuple (nodes_dict, links_list) where: + nodes_dict is a mapping of node_name -> Node instance. + links_list is a list of Link instances. + """ nodes = self._build_nodes() links = self._build_links(nodes) return nodes, links def _build_nodes(self): + """ + Internal method to build Node objects. + + :return: Dictionary of node_name -> Node instances. + """ nodes_from_clab = self.containerlab_data["topology"]["nodes"] node_width = self.styles.get("node_width", 75) @@ -51,6 +72,12 @@ def _build_nodes(self): return nodes def _build_links(self, nodes): + """ + Internal method to build Link objects and attach them to nodes. + + :param nodes: Dictionary of node_name -> Node instances. + :return: List of Link instances. + """ links_from_clab = [] for link in self.containerlab_data["topology"].get("links", []): endpoints = link.get("endpoints") diff --git a/core/data/topology_loader.py b/core/data/topology_loader.py index 62c223f..2bb777e 100644 --- a/core/data/topology_loader.py +++ b/core/data/topology_loader.py @@ -1,21 +1,34 @@ import yaml import sys -import os +import logging + +logger = logging.getLogger(__name__) + +class TopologyLoaderError(Exception): + """Raised when loading the topology fails.""" class TopologyLoader: - def load(self, input_file): + """ + Loads containerlab topology data from a YAML file. + """ + def load(self, input_file: str) -> dict: """ Load the containerlab YAML topology file and return its contents as a dictionary. + + :param input_file: Path to the containerlab YAML file. + :return: Parsed containerlab topology data. """ + logger.debug(f"Loading topology from file: {input_file}") try: with open(input_file, "r") as file: containerlab_data = yaml.safe_load(file) + logger.debug("Topology successfully loaded.") return containerlab_data except FileNotFoundError: error_message = f"Error: The specified clab file '{input_file}' does not exist." - print(error_message) - sys.exit(1) + logger.error(error_message) + raise TopologyLoaderError(error_message) except Exception as e: error_message = f"An error occurred while loading the config: {e}" - print(error_message) - sys.exit(1) + logger.error(error_message) + raise TopologyLoaderError(error_message) diff --git a/core/diagram/diagram_builder.py b/core/diagram/diagram_builder.py index 2310dbc..870185c 100644 --- a/core/diagram/diagram_builder.py +++ b/core/diagram/diagram_builder.py @@ -2,7 +2,17 @@ import random class DiagramBuilder: + """ + Builds diagram elements such as nodes, ports, and links into the Draw.io diagram. + """ def add_ports(self, diagram, styles, verbose=True): + """ + Add ports to the diagram nodes based on their links and layout. + + :param diagram: CustomDrawioDiagram instance. + :param styles: Dictionary of style parameters. + :param verbose: Whether to print debugging information. + """ nodes = diagram.nodes for node in nodes.values(): @@ -237,6 +247,12 @@ def add_ports(self, diagram, styles, verbose=True): ) def add_links(self, diagram, styles): + """ + Add links to the diagram, including link labels and styling. + + :param diagram: CustomDrawioDiagram instance. + :param styles: Dictionary of style parameters. + """ nodes = diagram.nodes global_seen_links = set() @@ -340,6 +356,13 @@ def add_links(self, diagram, styles): ) def add_nodes(self, diagram, nodes, styles): + """ + Add nodes to the diagram with their computed positions and styling. + + :param diagram: CustomDrawioDiagram instance. + :param nodes: Dictionary of node_name -> Node instances. + :param styles: Dictionary of style parameters. + """ base_style = styles["base_style"] custom_styles = styles["custom_styles"] icon_to_group_mapping = styles["icon_to_group_mapping"] diff --git a/core/interactivity/interactive_manager.py b/core/interactivity/interactive_manager.py index 2cc80f3..78ae894 100644 --- a/core/interactivity/interactive_manager.py +++ b/core/interactivity/interactive_manager.py @@ -3,17 +3,31 @@ import re class InteractiveManager: + """ + Manages interactive mode for assigning graph-levels and graph-icons via a CLI interface. + """ def run_interactive_mode( self, - nodes, - icon_to_group_mapping, - containerlab_data, - output_file, + nodes: dict, + icon_to_group_mapping: dict, + containerlab_data: dict, + output_file: str, processor, - prefix, - lab_name, - ): - # Original logic from interactive_mode() copied here without changes: + prefix: str, + lab_name: str, + ) -> dict: + """ + Run the interactive mode dialogs to set graph-levels and icons for nodes. + + :param nodes: Dictionary of node_name -> Node instances. + :param icon_to_group_mapping: Mapping from icon names to style groups. + :param containerlab_data: Parsed containerlab topology data. + :param output_file: Path to the output containerlab YAML file. + :param processor: YAMLProcessor instance for saving updated YAML. + :param prefix: Node name prefix. + :param lab_name: Lab name string. + :return: Summary dictionary of the chosen configuration. + """ previous_summary = {"Levels": {}, "Icons": {}} for node_name, node in nodes.items(): try: diff --git a/core/layout/horizontal_layout.py b/core/layout/horizontal_layout.py index 2f33e77..f8c19d4 100644 --- a/core/layout/horizontal_layout.py +++ b/core/layout/horizontal_layout.py @@ -2,7 +2,16 @@ from collections import defaultdict class HorizontalLayout(LayoutManager): - def apply(self, diagram, verbose=False): + """ + Applies a horizontal layout strategy to arrange nodes in the diagram. + """ + def apply(self, diagram, verbose=False) -> None: + """ + Apply a horizontal layout to the diagram nodes. + + :param diagram: CustomDrawioDiagram instance with nodes. + :param verbose: Whether to print debugging information. + """ self.diagram = diagram self.verbose = verbose self._calculate_positions() diff --git a/core/layout/vertical_layout.py b/core/layout/vertical_layout.py index c662305..95ff2a1 100644 --- a/core/layout/vertical_layout.py +++ b/core/layout/vertical_layout.py @@ -2,7 +2,16 @@ from collections import defaultdict class VerticalLayout(LayoutManager): - def apply(self, diagram, verbose=False): + """ + Applies a vertical layout strategy to arrange nodes in the diagram. + """ + def apply(self, diagram, verbose=False) -> None: + """ + Apply a vertical layout to the diagram nodes. + + :param diagram: CustomDrawioDiagram instance with nodes. + :param verbose: Whether to print debugging information. + """ self.diagram = diagram self.verbose = verbose self._calculate_positions() diff --git a/core/logging_config.py b/core/logging_config.py new file mode 100644 index 0000000..50a6344 --- /dev/null +++ b/core/logging_config.py @@ -0,0 +1,12 @@ +import logging + +def configure_logging(level=logging.INFO): + """ + Configure the logging settings for the application. + Adjust the format and level as desired. + """ + logging.basicConfig( + level=level, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S" + ) diff --git a/drawio2clab.py b/drawio2clab.py index e7b499a..0bd8b7a 100644 --- a/drawio2clab.py +++ b/drawio2clab.py @@ -5,7 +5,22 @@ from core.drawio.drawio_parser import DrawioParser from core.drawio.converter import Drawio2ClabConverter -def main(input_file, output_file, style="flow", diagram_name=None, default_kind="nokia_srlinux"): +def main( + input_file: str, + output_file: str, + style: str="flow", + diagram_name: str=None, + default_kind: str="nokia_srlinux" +) -> None: + """ + Main function to convert a .drawio file to a Containerlab YAML file. + + :param input_file: Path to the .drawio XML file. + :param output_file: Output YAML file path. + :param style: YAML style ("block" or "flow"). + :param diagram_name: Name of the diagram to parse within the .drawio file. + :param default_kind: Default kind for nodes if not specified in the diagram. + """ parser = DrawioParser(input_file, diagram_name) mxGraphModel_root = parser.parse_xml()