Init: mediaserver

lib/python3.10/site-packages/ansibleplaybookgrapher/graphbuilder.py

@@ -0,0 +1,480 @@
+# Copyright (C) 2022 Mohamed El Mouctar HAIDARA
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+from typing import Dict, Optional, Tuple, List, Set
+
+from ansible.utils.display import Display
+from graphviz import Digraph
+
+from ansibleplaybookgrapher import PlaybookParser
+from ansibleplaybookgrapher.graph import (
+    PlaybookNode,
+    RoleNode,
+    BlockNode,
+    Node,
+    PlayNode,
+)
+from ansibleplaybookgrapher.utils import get_play_colors, merge_dicts
+
+display = Display()
+
+# The supported protocol handlers to open roles and tasks from the viewer
+OPEN_PROTOCOL_HANDLERS = {
+    "default": {"folder": "{path}", "file": "{path}"},
+    # https://code.visualstudio.com/docs/editor/command-line#_opening-vs-code-with-urls
+    "vscode": {
+        "folder": "vscode://file/{path}",
+        "file": "vscode://file/{path}:{line}:{column}",
+    },
+    # For custom, the formats need to be provided
+    "custom": {},
+}
+
+
+class Grapher:
+    def __init__(self, playbook_filenames: List[str]):
+        """
+
+        :param playbook_filenames: List of playbooks to graph
+        """
+        self.playbook_filenames = playbook_filenames
+        # Colors assigned to plays
+
+        self.plays_color = {}
+        # The usage of the roles in all playbooks
+        self.roles_usage: Dict["RoleNode", List[str]] = {}
+
+        # The parsed playbooks
+        self.playbook_nodes: List[PlaybookNode] = []
+
+    def parse(
+        self,
+        include_role_tasks: bool = False,
+        tags: List[str] = None,
+        skip_tags: List[str] = None,
+        group_roles_by_name: bool = False,
+    ):
+        """
+        Parses all the provided playbooks
+        :param include_role_tasks: Should we include the role tasks
+        :param tags: Only add plays and tasks tagged with these values
+        :param skip_tags: Only add plays and tasks whose tags do not match these values
+        :param group_roles_by_name: Group roles by name instead of considering them as separate nodes with different IDs
+        :return:
+        """
+
+        for playbook_file in self.playbook_filenames:
+            display.display(f"Parsing playbook {playbook_file}")
+            parser = PlaybookParser(
+                tags=tags,
+                skip_tags=skip_tags,
+                playbook_filename=playbook_file,
+                include_role_tasks=include_role_tasks,
+                group_roles_by_name=group_roles_by_name,
+            )
+            playbook_node = parser.parse()
+            self.playbook_nodes.append(playbook_node)
+
+            # Setting colors for play
+            for play in playbook_node.plays:
+                # TODO: find a way to create visual distance between the generated colors
+                #   https://stackoverflow.com/questions/9018016/how-to-compare-two-colors-for-similarity-difference
+                self.plays_color[play] = get_play_colors(play.id)
+
+            # Update the usage of the roles
+            self.roles_usage = merge_dicts(
+                self.roles_usage, playbook_node.roles_usage()
+            )
+
+    def graph(
+        self,
+        open_protocol_handler: str,
+        open_protocol_custom_formats: Dict[str, str] = None,
+    ) -> Digraph:
+        """
+        Generate the digraph graph
+        :param open_protocol_handler
+        :param open_protocol_custom_formats
+        :return:
+        """
+        digraph = Digraph(
+            format="svg",
+            graph_attr=GraphvizGraphBuilder.DEFAULT_GRAPH_ATTR,
+            edge_attr=GraphvizGraphBuilder.DEFAULT_EDGE_ATTR,
+        )
+        # Map of the rules that have been built so far for all playbooks
+        roles_built = {}
+        for p in self.playbook_nodes:
+            builder = GraphvizGraphBuilder(
+                p,
+                digraph=digraph,
+                roles_usage=self.roles_usage,
+                roles_built=roles_built,
+                play_colors=self.plays_color,
+                open_protocol_handler=open_protocol_handler,
+                open_protocol_custom_formats=open_protocol_custom_formats,
+            )
+            builder.build_graphviz_graph()
+            roles_built.update(builder.roles_built)
+
+        return digraph
+
+
+class GraphvizGraphBuilder:
+    """
+    Build the graphviz graph
+    """
+
+    DEFAULT_EDGE_ATTR = {"sep": "10", "esep": "5"}
+    DEFAULT_GRAPH_ATTR = {
+        "ratio": "fill",
+        "rankdir": "LR",
+        "concentrate": "true",
+        "ordering": "in",
+    }
+
+    def __init__(
+        self,
+        playbook_node: PlaybookNode,
+        open_protocol_handler: str,
+        digraph: Digraph,
+        play_colors: Dict[PlayNode, Tuple[str, str]],
+        roles_usage: Dict[RoleNode, List[Node]] = None,
+        roles_built: Dict = None,
+        open_protocol_custom_formats: Dict[str, str] = None,
+    ):
+        """
+
+        :param playbook_node: Playbook parsed node
+        :param open_protocol_handler: The protocol handler name to use
+        :param digraph: Graphviz graph into which build the graph
+        :param open_protocol_custom_formats: The custom formats to use when the protocol handler is set to custom
+        """
+        self.playbook_node = playbook_node
+        self.roles_usage = roles_usage or playbook_node.roles_usage()
+        self.play_colors = play_colors
+        # A map containing the roles that have been built so far
+        self.roles_built = roles_built or {}
+
+        self.open_protocol_handler = open_protocol_handler
+        # Merge the two dicts
+        formats = {**OPEN_PROTOCOL_HANDLERS, **{"custom": open_protocol_custom_formats}}
+        self.open_protocol_formats = formats[self.open_protocol_handler]
+
+        self.digraph = digraph
+
+    def build_node(
+        self,
+        graph: Digraph,
+        counter: int,
+        source: Node,
+        destination: Node,
+        color: str,
+        shape: str = "octagon",
+        **kwargs,
+    ):
+        """
+        Render a generic node in the graph
+        :param graph: The graph to render the node to
+        :param source: The source node
+        :param destination: The RoleNode to render
+        :param color: The color to apply
+        :param counter: The counter for this node
+        :param shape: the default shape of the node
+        :return:
+        """
+
+        node_label_prefix = kwargs.get("node_label_prefix", "")
+
+        if isinstance(destination, BlockNode):
+            self.build_block(
+                graph,
+                counter,
+                source=source,
+                destination=destination,
+                color=color,
+            )
+        elif isinstance(destination, RoleNode):
+            self.build_role(
+                graph,
+                counter,
+                source=source,
+                destination=destination,
+                color=color,
+            )
+        else:
+            # Here we have a TaskNode
+            edge_label = f"{counter} {destination.when}"
+            # Task node
+            graph.node(
+                destination.id,
+                label=node_label_prefix + destination.name,
+                shape=shape,
+                id=destination.id,
+                tooltip=destination.name,
+                color=color,
+                URL=self.get_node_url(destination, "file"),
+            )
+            # Edge from parent to task
+            graph.edge(
+                source.id,
+                destination.id,
+                label=edge_label,
+                color=color,
+                fontcolor=color,
+                id=f"edge_{counter}_{source.id}_{destination.id}",
+                tooltip=edge_label,
+                labeltooltip=edge_label,
+            )
+
+    def build_block(
+        self,
+        graph: Digraph,
+        counter: int,
+        source: Node,
+        destination: BlockNode,
+        color: str,
+        **kwargs,
+    ):
+        """
+        Render a block in the graph.
+        A BlockNode is a special node: a cluster is created instead of a normal node.
+        :param graph: The graph to render the block into
+        :param counter: The counter for this block in the graph
+        :param source: The source node
+        :param destination: The BlockNode to render
+        :param color: The color to apply
+        :param kwargs:
+        :return:
+        """
+        edge_label = f"{counter}"
+
+        # Edge from parent to the block node inside the cluster
+        graph.edge(
+            source.id,
+            destination.id,
+            label=edge_label,
+            color=color,
+            fontcolor=color,
+            tooltip=edge_label,
+            id=f"edge_{counter}_{source.id}_{destination.id}",
+            labeltooltip=edge_label,
+        )
+
+        # BlockNode is a special node: a cluster is created instead of a normal node
+        with graph.subgraph(name=f"cluster_{destination.id}") as cluster_block_subgraph:
+            # block node
+            cluster_block_subgraph.node(
+                destination.id,
+                label=f"[block] {destination.name}",
+                shape="box",
+                id=destination.id,
+                tooltip=destination.name,
+                color=color,
+                labeltooltip=destination.name,
+                URL=self.get_node_url(destination, "file"),
+            )
+
+            # The reverse here is a little hack due to how graphviz render nodes inside a cluster by reversing them.
+            #  Don't really know why for the moment neither if there is an attribute to change that.
+            for b_counter, task in enumerate(reversed(destination.tasks)):
+                self.build_node(
+                    cluster_block_subgraph,
+                    source=destination,
+                    destination=task,
+                    counter=len(destination.tasks) - b_counter,
+                    color=color,
+                )
+
+    def build_role(
+        self,
+        graph: Digraph,
+        counter: int,
+        source: Node,
+        destination: RoleNode,
+        color: str,
+        **kwargs,
+    ):
+        """
+        Render a role in the graph
+        :param graph: The graph to render the role into
+        :param counter: The counter for this role in the graph
+        :param source: The source node
+        :param destination: The RoleNode to render
+        :param color: The color to apply
+        :param kwargs:
+        :return:
+        """
+
+        if destination.include_role:  # For include_role, we point to a file
+            url = self.get_node_url(destination, "file")
+        else:  # For normal role invocation, we point to the folder
+            url = self.get_node_url(destination, "folder")
+
+        role_edge_label = f"{counter} {destination.when}"
+
+        # from parent to the role node
+        graph.edge(
+            source.id,
+            destination.id,
+            label=role_edge_label,
+            color=color,
+            fontcolor=color,
+            id=f"edge_{counter}_{source.id}_{destination.id}",
+            tooltip=role_edge_label,
+            labeltooltip=role_edge_label,
+        )
+
+        # check if we already built this role
+        role_to_render = self.roles_built.get(destination.id, None)
+        if role_to_render is None:
+            # Merge the colors for each play where this role is used
+            role_plays = self.roles_usage[destination]
+            # Graphviz support providing multiple colors separated by :
+            if len(role_plays) > 1:
+                # If the role is used in multiple plays, we take black as the default color
+                role_color = "black"
+            else:
+                colors = list(map(self.play_colors.get, role_plays))[0]
+                role_color = colors[0]
+
+            self.roles_built[destination.id] = destination
+
+            with graph.subgraph(name=destination.name, node_attr={}) as role_subgraph:
+                role_subgraph.node(
+                    destination.id,
+                    id=destination.id,
+                    label=f"[role] {destination.name}",
+                    tooltip=destination.name,
+                    color=color,
+                    URL=url,
+                )
+                # role tasks
+                for role_task_counter, role_task in enumerate(destination.tasks, 1):
+                    self.build_node(
+                        role_subgraph,
+                        source=destination,
+                        destination=role_task,
+                        counter=role_task_counter,
+                        color=role_color,
+                    )
+
+    def build_graphviz_graph(self):
+        """
+        Convert the PlaybookNode to the graphviz dot format
+        :return:
+        """
+        display.vvv(f"Converting the graph to the dot format for graphviz")
+        # root node
+        self.digraph.node(
+            self.playbook_node.name,
+            style="dotted",
+            id=self.playbook_node.id,
+            URL=self.get_node_url(self.playbook_node, "file"),
+        )
+
+        for play_counter, play in enumerate(self.playbook_node.plays, 1):
+            with self.digraph.subgraph(name=play.name) as play_subgraph:
+                color, play_font_color = self.play_colors[play]
+                play_tooltip = (
+                    ",".join(play.hosts) if len(play.hosts) > 0 else play.name
+                )
+
+                # play node
+                play_subgraph.node(
+                    play.id,
+                    id=play.id,
+                    label=play.name,
+                    style="filled",
+                    shape="box",
+                    color=color,
+                    fontcolor=play_font_color,
+                    tooltip=play_tooltip,
+                    URL=self.get_node_url(play, "file"),
+                )
+
+                # edge from root node to play
+                playbook_to_play_label = f"{play_counter} {play.name}"
+                self.digraph.edge(
+                    self.playbook_node.name,
+                    play.id,
+                    id=f"edge_{self.playbook_node.id}_{play.id}",
+                    label=playbook_to_play_label,
+                    color=color,
+                    fontcolor=color,
+                    tooltip=playbook_to_play_label,
+                    labeltooltip=playbook_to_play_label,
+                )
+
+                # pre_tasks
+                for pre_task_counter, pre_task in enumerate(play.pre_tasks, 1):
+                    self.build_node(
+                        play_subgraph,
+                        counter=pre_task_counter,
+                        source=play,
+                        destination=pre_task,
+                        color=color,
+                        node_label_prefix="[pre_task] ",
+                    )
+
+                # roles
+                for role_counter, role in enumerate(play.roles, 1):
+                    self.build_role(
+                        play_subgraph,
+                        source=play,
+                        destination=role,
+                        counter=role_counter + len(play.pre_tasks),
+                        color=color,
+                    )
+
+                # tasks
+                for task_counter, task in enumerate(play.tasks, 1):
+                    self.build_node(
+                        play_subgraph,
+                        source=play,
+                        destination=task,
+                        counter=len(play.pre_tasks) + len(play.roles) + task_counter,
+                        color=color,
+                        node_label_prefix="[task] ",
+                    )
+
+                # post_tasks
+                for post_task_counter, post_task in enumerate(play.post_tasks, 1):
+                    self.build_node(
+                        play_subgraph,
+                        source=play,
+                        destination=post_task,
+                        counter=len(play.pre_tasks)
+                        + len(play.roles)
+                        + len(play.tasks)
+                        + post_task_counter,
+                        color=color,
+                        node_label_prefix="[post_task] ",
+                    )
+
+    def get_node_url(self, node: Node, node_type: str) -> Optional[str]:
+        """
+        Get the node url based on the chosen protocol
+        :param node_type: task or role
+        :param node: the node to get the url for
+        :return:
+        """
+        if node.path:
+            url = self.open_protocol_formats[node_type].format(
+                path=node.path, line=node.line, column=node.column
+            )
+            display.vvvv(f"Open protocol URL for node {node}: {url}")
+            return url
+
+        return None