Use a dataclass for position + simplification

ansibleplaybookgrapher/graph_model.py

@@ -14,6 +14,7 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 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:

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

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

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))

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]