diff --git a/compose_viz/cli.py b/compose_viz/cli.py index 00c1ddd..8e4ef9b 100644 --- a/compose_viz/cli.py +++ b/compose_viz/cli.py @@ -3,7 +3,7 @@ import typer from typing import Optional from compose_viz import __app_name__, __version__ from compose_viz.parser import Parser - +from compose_viz.graph import Graph class VisualizationFormats(str, Enum): png = "PNG" @@ -54,6 +54,8 @@ def compose_viz( if compose: typer.echo(f"Successfully parsed {input_path}") + Graph(compose, output_path).render(format) + raise typer.Exit() diff --git a/compose_viz/graph.py b/compose_viz/graph.py new file mode 100644 index 0000000..9f8a7fa --- /dev/null +++ b/compose_viz/graph.py @@ -0,0 +1,72 @@ +import graphviz + +from compose_viz.compose import Compose + + +def apply_vertex_style(type) -> dict: + style = { + 'service': { + 'shape': 'component', + }, + 'volume': { + 'shape': 'folder', + }, + 'network': { + 'shape': 'pentagon', + }, + 'port': { + 'shape': 'circle', + }, + } + + return style[type] + + +def apply_edge_style(type) -> dict: + style = { + 'ports': { + 'style': 'solid', + }, + 'links': { + 'style': 'solid', + }, + 'volumes': { + 'style': 'dashed', + }, + 'depends_on': { + 'style': 'dotted', + } + } + + return style[type] + + +class Graph: + def __init__(self, compose: Compose, filename: str) -> None: + self.dot = graphviz.Digraph() + self.dot.attr('graph', background='#ffffff', pad='0.5', ratio='fill') + self.compose = compose + self.filename = filename + + def add_vertex(self, name: str, type: str) -> None: + self.dot.node(name, **apply_vertex_style(type)) + + def add_edge(self, head: str, tail: str, type: str) -> None: + self.dot.edge(head, tail, **apply_edge_style(type)) + + def render(self, format: str, cleanup: bool = True) -> None: + for service in self.compose.services: + self.add_vertex(service.name, 'service') + for network in service.networks: + self.add_vertex("net#" + network, 'network') + self.add_edge(service.name, "net#" + network, 'links') + for volume in service.volumes: + self.add_vertex(volume, 'volume') + self.add_edge(service.name, volume, 'links') + for port in service.ports: + self.add_vertex(port, 'port') + self.add_edge(service.name, port, 'ports') + for depends_on in service.depends_on: + self.dot.edge(depends_on, service.name, 'depends_on') + + self.dot.render(outfile=self.filename, format=format, cleanup=cleanup) diff --git a/tests/test_cli.py b/tests/test_cli.py index 3f427b5..bc0841e 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,3 +1,5 @@ +import os + from typer.testing import CliRunner from compose_viz import cli @@ -11,3 +13,4 @@ def test_cli(): assert result.exit_code == 0 assert f"Successfully parsed {input_path}\n" in result.stdout + assert os.path.exists("compose-viz.png")