From 8909b25a4f8ec9e257752a12259dbf6dd34e9bd6 Mon Sep 17 00:00:00 2001 From: Mohamed El Mouctar HAIDARA Date: Wed, 17 Jul 2024 11:29:51 +0000 Subject: [PATCH] Use a dataclass for position + simplification --- ansibleplaybookgrapher/graph_model.py | 67 +++++++++++++------ ansibleplaybookgrapher/renderer/__init__.py | 12 ++-- .../renderer/graphviz/__init__.py | 12 ++-- tests/test_graph_model.py | 10 ++- tests/test_parser.py | 24 +++---- 5 files changed, 77 insertions(+), 48 deletions(-) diff --git a/ansibleplaybookgrapher/graph_model.py b/ansibleplaybookgrapher/graph_model.py index 186bbfdb..18980463 100644 --- a/ansibleplaybookgrapher/graph_model.py +++ b/ansibleplaybookgrapher/graph_model.py @@ -14,6 +14,7 @@ # along with this program. If not, see . import os from collections import defaultdict +from dataclasses import dataclass, asdict from typing import Dict, List, Set, Tuple, Optional from ansibleplaybookgrapher.utils import generate_id, get_play_colors @@ -35,6 +36,24 @@ def has_loop(self) -> bool: return self.raw_object.loop is not None +@dataclass +class NodeLocation: + """ + The node location on the disk. The location can be a folder (for roles) or a specific line and column inside a file + """ + + type: str # file or folder + path: Optional[str] = None + line: Optional[int] = None + column: Optional[int] = None + + def __post_init__(self): + if self.type not in ["folder", "file", None]: + raise ValueError( + f"Type '{self.type}' not supported. Valid values: file, folder." + ) + + class Node: """ A node in the graph. Everything of the final graph is a node: playbook, plays, tasks and roles. @@ -64,20 +83,23 @@ def __init__( self.when = when self.raw_object = raw_object - # Get the node position in the parsed files. Format: (path,line,column) - self.path = self.line = self.column = None - self.set_position() + self.location = None + self.set_location() # The index of this node in the parent node if it has one (starting from 1) self.index: Optional[int] = index - def set_position(self): + def set_location(self): """ - Set the path of this based on the raw object. Not all objects have path + Set the location of this node based on the raw object. Not all objects have path :return: """ if self.raw_object and self.raw_object.get_ds(): - self.path, self.line, self.column = self.raw_object.get_ds().ansible_pos + path, line, column = self.raw_object.get_ds().ansible_pos + # By default, it's a file + self.location = NodeLocation( + type="file", path=path, line=line, column=column + ) def get_first_parent_matching_type(self, node_type: type) -> type: """ @@ -112,19 +134,21 @@ def to_dict(self, **kwargs) -> Dict: back. :return: """ - return { + data = { "type": type(self).__name__, "id": self.id, - "position": { - "path": self.path, - "line": self.line, - "column": self.column, - }, "name": self.name, "when": self.when, "index": self.index, } + if self.location is not None: + data["location"] = asdict(self.location) + else: + data["location"] = None + + return data + class CompositeNode(Node): """ @@ -351,15 +375,15 @@ def __init__( supported_compositions=["plays"], ) - def set_position(self): + def set_location(self): """ Playbooks only have path as position :return: """ # Since the playbook is the whole file, the set the position as the beginning of the file - self.path = os.path.join(os.getcwd(), self.name) - self.line = 1 - self.column = 1 + self.location = NodeLocation( + type="file", path=os.path.join(os.getcwd(), self.name), line=1, column=1 + ) def plays( self, exclude_empty: bool = False, exclude_without_roles: bool = False @@ -580,17 +604,18 @@ def __init__( index=index, ) - def set_position(self): + def set_location(self): """ Retrieve the position depending on whether it's an include_role or not :return: """ if self.raw_object and not self.include_role: - # If it's not an include_role, we take the role path which the path to the folder where the role is located - # on the disk - self.path = self.raw_object._role_path + # If it's not an include_role, we take the role path which is the path to the folder where the role + # is located on the disk. + self.location = NodeLocation(type="folder", path=self.raw_object._role_path) + else: - super().set_position() + super().set_location() def has_loop(self) -> bool: if not self.include_role: diff --git a/ansibleplaybookgrapher/renderer/__init__.py b/ansibleplaybookgrapher/renderer/__init__.py index ae09cc51..4e41912b 100644 --- a/ansibleplaybookgrapher/renderer/__init__.py +++ b/ansibleplaybookgrapher/renderer/__init__.py @@ -239,19 +239,19 @@ def build_block(self, block_node: BlockNode, color: str, fontcolor: str, **kwarg """ pass - def get_node_url(self, node: Node, node_type: str) -> Optional[str]: + def get_node_url(self, node: Node) -> Optional[str]: """ Get the node url based on the chosen open protocol. - :param node_type: task or role + :param node: the node to get the url for :return: """ - if node.path: + if node.location.path: remove_from_path = self.open_protocol_formats.get("remove_from_path", "") - path = node.path.replace(remove_from_path, "") + path = node.location.path.replace(remove_from_path, "") - url = self.open_protocol_formats[node_type].format( - path=path, line=node.line, column=node.column + url = self.open_protocol_formats[node.location.type].format( + path=path, line=node.location.line, column=node.location.column ) display.vvvv(f"Open protocol URL for node {node}: {url}") return url diff --git a/ansibleplaybookgrapher/renderer/graphviz/__init__.py b/ansibleplaybookgrapher/renderer/graphviz/__init__.py index b410c288..8f31b48d 100644 --- a/ansibleplaybookgrapher/renderer/graphviz/__init__.py +++ b/ansibleplaybookgrapher/renderer/graphviz/__init__.py @@ -162,7 +162,7 @@ def build_task(self, task_node: TaskNode, color: str, fontcolor: str, **kwargs): id=task_node.id, tooltip=task_node.name, color=color, - URL=self.get_node_url(task_node, "file"), + URL=self.get_node_url(task_node), ) # Edge from parent to task @@ -212,7 +212,7 @@ def build_block(self, block_node: BlockNode, color: str, fontcolor: str, **kwarg color=color, fontcolor=fontcolor, labeltooltip=block_node.name, - URL=self.get_node_url(block_node, "file"), + URL=self.get_node_url(block_node), ) # The reverse here is a little hack due to how graphviz render nodes inside a cluster by reversing them. @@ -253,9 +253,9 @@ def build_role(self, role_node: RoleNode, color: str, fontcolor: str, **kwargs): self.roles_built.add(role_node) if role_node.include_role: # For include_role, we point to a file - url = self.get_node_url(role_node, "file") + url = self.get_node_url(role_node) else: # For normal role invocation, we point to the folder - url = self.get_node_url(role_node, "folder") + url = self.get_node_url(role_node) plays_using_this_role = self.roles_usage[role_node] if len(plays_using_this_role) > 1: @@ -301,7 +301,7 @@ def build_playbook( label=self.playbook_node.name, style="dotted", id=self.playbook_node.id, - URL=self.get_node_url(self.playbook_node, "file"), + URL=self.get_node_url(self.playbook_node), ) for play in self.playbook_node.plays( @@ -337,7 +337,7 @@ def build_play(self, play_node: PlayNode, **kwargs): color=color, fontcolor=play_font_color, tooltip=play_tooltip, - URL=self.get_node_url(play_node, "file"), + URL=self.get_node_url(play_node), ) # from playbook to play diff --git a/tests/test_graph_model.py b/tests/test_graph_model.py index 61e916c5..baf4f5de 100644 --- a/tests/test_graph_model.py +++ b/tests/test_graph_model.py @@ -1,3 +1,5 @@ +import json + from ansibleplaybookgrapher.graph_model import ( RoleNode, TaskNode, @@ -124,9 +126,9 @@ def test_to_dict(): dict_rep = playbook.to_dict(exclude_empty_plays=True) assert dict_rep["type"] == "PlaybookNode" - assert dict_rep["position"]["path"] is not None - assert dict_rep["position"]["line"] is not None - assert dict_rep["position"]["column"] is not None + assert dict_rep["location"]["path"] is not None + assert dict_rep["location"]["line"] is not None + assert dict_rep["location"]["column"] is not None assert len(dict_rep["plays"]) == 1 assert dict_rep["plays"][0]["type"] == "PlayNode" @@ -136,3 +138,5 @@ def test_to_dict(): assert dict_rep["plays"][0]["tasks"][0]["name"] == "block 1" assert dict_rep["plays"][0]["tasks"][0]["index"] == 1 assert dict_rep["plays"][0]["tasks"][0]["type"] == "BlockNode" + + print(json.dumps(dict_rep, indent=4)) diff --git a/tests/test_parser.py b/tests/test_parser.py index 70929170..e6d09879 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -49,16 +49,16 @@ def test_example_parsing(grapher_cli: PlaybookGrapherCLI, display: Display): parser = PlaybookParser(grapher_cli.options.playbook_filenames[0]) playbook_node = parser.parse() assert len(playbook_node.plays()) == 1 - assert playbook_node.path == os.path.join(FIXTURES_PATH, "example.yml") - assert playbook_node.line == 1 - assert playbook_node.column == 1 + assert playbook_node.location.path == os.path.join(FIXTURES_PATH, "example.yml") + assert playbook_node.location.line == 1 + assert playbook_node.location.column == 1 assert ( playbook_node.index is None ), "The index of the playbook should be None (it has no parent)" play_node = playbook_node.plays()[0] - assert play_node.path == os.path.join(FIXTURES_PATH, "example.yml") - assert play_node.line == 2 + assert play_node.location.path == os.path.join(FIXTURES_PATH, "example.yml") + assert play_node.location.line == 2 assert play_node.index == 1 pre_tasks = play_node.pre_tasks @@ -99,9 +99,9 @@ def test_with_roles_parsing(grapher_cli: PlaybookGrapherCLI): fake_role = play_node.roles[0] assert isinstance(fake_role, RoleNode) assert not fake_role.include_role - assert fake_role.path == os.path.join(FIXTURES_PATH, "roles", "fake_role") - assert fake_role.line is None - assert fake_role.column is None + assert fake_role.location.path == os.path.join(FIXTURES_PATH, "roles", "fake_role") + assert fake_role.location.line is None + assert fake_role.location.column is None assert fake_role.index == 3 for task_counter, task in enumerate(fake_role.tasks): @@ -144,8 +144,8 @@ def test_include_role_parsing(grapher_cli: PlaybookGrapherCLI, capsys): include_role_1 = block_include_role.tasks[0] assert isinstance(include_role_1, RoleNode) assert include_role_1.include_role - assert include_role_1.path == os.path.join(FIXTURES_PATH, "include_role.yml") - assert include_role_1.line == 10, "The first include role should be at line 9" + assert include_role_1.location.path == os.path.join(FIXTURES_PATH, "include_role.yml") + assert include_role_1.location.line == 10, "The first include role should be at line 9" assert ( len(include_role_1.tasks) == 0 ), "We don't support adding tasks from include_role with loop" @@ -223,8 +223,8 @@ def test_block_parsing(grapher_cli: PlaybookGrapherCLI): assert isinstance( pre_task_block, BlockNode ), "The second edge should have a BlockNode as destination" - assert pre_task_block.path == os.path.join(FIXTURES_PATH, "with_block.yml") - assert pre_task_block.line == 7 + assert pre_task_block.location.path == os.path.join(FIXTURES_PATH, "with_block.yml") + assert pre_task_block.location.line == 7 # Check tasks task_1 = tasks[0]