coco-framework 0.1.0__py3-none-any.whl

coco_core/runnables/graph_ascii.py

@@ -0,0 +1,366 @@
+"""Draws DAG in ASCII.
+
+Adapted from https://github.com/iterative/dvc/blob/main/dvc/dagascii.py.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+from typing import TYPE_CHECKING, Any
+
+try:
+    from grandalf.graphs import Edge, Graph, Vertex  # type: ignore[import-untyped]
+    from grandalf.layouts import SugiyamaLayout  # type: ignore[import-untyped]
+    from grandalf.routing import route_with_lines  # type: ignore[import-untyped]
+
+    _HAS_GRANDALF = True
+except ImportError:
+    _HAS_GRANDALF = False
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+
+    from coco_core.runnables.graph import Edge as LangEdge
+
+
+class VertexViewer:
+    """VertexViewer class.
+
+    Class to define vertex box boundaries that will be accounted for during
+    graph building by grandalf.
+    """
+
+    HEIGHT = 3  # top and bottom box edges + text
+    """Height of the box."""
+
+    def __init__(self, name: str) -> None:
+        """Create a VertexViewer.
+
+        Args:
+            name: name of the vertex.
+        """
+        self._h = self.HEIGHT  # top and bottom box edges + text
+        self._w = len(name) + 2  # right and left bottom edges + text
+
+    @property
+    def h(self) -> int:
+        """Height of the box."""
+        return self._h
+
+    @property
+    def w(self) -> int:
+        """Width of the box."""
+        return self._w
+
+
+class AsciiCanvas:
+    """Class for drawing in ASCII."""
+
+    TIMEOUT = 10
+
+    def __init__(self, cols: int, lines: int) -> None:
+        """Create an ASCII canvas.
+
+        Args:
+            cols: number of columns in the canvas. Should be `> 1`.
+            lines: number of lines in the canvas. Should be `> 1`.
+
+        Raises:
+            ValueError: if canvas dimensions are invalid.
+        """
+        if cols <= 1 or lines <= 1:
+            msg = "Canvas dimensions should be > 1"
+            raise ValueError(msg)
+
+        self.cols = cols
+        self.lines = lines
+
+        self.canvas = [[" "] * cols for line in range(lines)]
+
+    def draw(self) -> str:
+        """Draws ASCII canvas on the screen.
+
+        Returns:
+            The ASCII canvas string.
+        """
+        lines = map("".join, self.canvas)
+        return os.linesep.join(lines)
+
+    def point(self, x: int, y: int, char: str) -> None:
+        """Create a point on ASCII canvas.
+
+        Args:
+            x: x coordinate. Should be `>= 0` and `<` number of columns in
+                the canvas.
+            y: y coordinate. Should be `>= 0` an `<` number of lines in the
+                canvas.
+            char: character to place in the specified point on the
+                canvas.
+
+        Raises:
+            ValueError: if char is not a single character or if
+                coordinates are out of bounds.
+        """
+        if len(char) != 1:
+            msg = "char should be a single character"
+            raise ValueError(msg)
+        if x >= self.cols or x < 0:
+            msg = "x should be >= 0 and < number of columns"
+            raise ValueError(msg)
+        if y >= self.lines or y < 0:
+            msg = "y should be >= 0 and < number of lines"
+            raise ValueError(msg)
+
+        self.canvas[y][x] = char
+
+    def line(self, x0: int, y0: int, x1: int, y1: int, char: str) -> None:
+        """Create a line on ASCII canvas.
+
+        Args:
+            x0: x coordinate where the line should start.
+            y0: y coordinate where the line should start.
+            x1: x coordinate where the line should end.
+            y1: y coordinate where the line should end.
+            char: character to draw the line with.
+        """
+        if x0 > x1:
+            x1, x0 = x0, x1
+            y1, y0 = y0, y1
+
+        dx = x1 - x0
+        dy = y1 - y0
+
+        if dx == 0 and dy == 0:
+            self.point(x0, y0, char)
+        elif abs(dx) >= abs(dy):
+            for x in range(x0, x1 + 1):
+                y = y0 if dx == 0 else y0 + round((x - x0) * dy / float(dx))
+                self.point(x, y, char)
+        elif y0 < y1:
+            for y in range(y0, y1 + 1):
+                x = x0 if dy == 0 else x0 + round((y - y0) * dx / float(dy))
+                self.point(x, y, char)
+        else:
+            for y in range(y1, y0 + 1):
+                x = x0 if dy == 0 else x1 + round((y - y1) * dx / float(dy))
+                self.point(x, y, char)
+
+    def text(self, x: int, y: int, text: str) -> None:
+        """Print a text on ASCII canvas.
+
+        Args:
+            x: x coordinate where the text should start.
+            y: y coordinate where the text should start.
+            text: string that should be printed.
+        """
+        for i, char in enumerate(text):
+            self.point(x + i, y, char)
+
+    def box(self, x0: int, y0: int, width: int, height: int) -> None:
+        """Create a box on ASCII canvas.
+
+        Args:
+            x0: x coordinate of the box corner.
+            y0: y coordinate of the box corner.
+            width: box width.
+            height: box height.
+
+        Raises:
+            ValueError: if box dimensions are invalid.
+        """
+        if width <= 1 or height <= 1:
+            msg = "Box dimensions should be > 1"
+            raise ValueError(msg)
+
+        width -= 1
+        height -= 1
+
+        for x in range(x0, x0 + width):
+            self.point(x, y0, "-")
+            self.point(x, y0 + height, "-")
+
+        for y in range(y0, y0 + height):
+            self.point(x0, y, "|")
+            self.point(x0 + width, y, "|")
+
+        self.point(x0, y0, "+")
+        self.point(x0 + width, y0, "+")
+        self.point(x0, y0 + height, "+")
+        self.point(x0 + width, y0 + height, "+")
+
+
+class _EdgeViewer:
+    def __init__(self) -> None:
+        self.pts: list[tuple[float]] = []
+
+    def setpath(self, pts: list[tuple[float]]) -> None:
+        self.pts = pts
+
+
+def _build_sugiyama_layout(
+    vertices: Mapping[str, str], edges: Sequence[LangEdge]
+) -> Any:
+    if not _HAS_GRANDALF:
+        msg = "Install grandalf to draw graphs: `pip install grandalf`."
+        raise ImportError(msg)
+
+    #
+    # Just a reminder about naming conventions:
+    # +------------X
+    # |
+    # |
+    # |
+    # |
+    # Y
+    #
+
+    vertices_ = {id_: Vertex(f" {data} ") for id_, data in vertices.items()}
+    edges_ = [Edge(vertices_[s], vertices_[e], data=cond) for s, e, _, cond in edges]
+    vertices_list = vertices_.values()
+    graph = Graph(vertices_list, edges_)
+
+    for vertex in vertices_list:
+        vertex.view = VertexViewer(vertex.data)
+
+    # NOTE: determine min box length to create the best layout
+    minw = min(v.view.w for v in vertices_list)
+
+    for edge in edges_:
+        edge.view = _EdgeViewer()
+
+    sug = SugiyamaLayout(graph.C[0])
+    graph = graph.C[0]
+    roots = list(filter(lambda x: len(x.e_in()) == 0, graph.sV))
+
+    sug.init_all(roots=roots, optimize=True)
+
+    sug.yspace = VertexViewer.HEIGHT
+    sug.xspace = minw
+    sug.route_edge = route_with_lines
+
+    sug.draw()
+
+    return sug
+
+
+def draw_ascii(vertices: Mapping[str, str], edges: Sequence[LangEdge]) -> str:
+    """Build a DAG and draw it in ASCII.
+
+    Args:
+        vertices: list of graph vertices.
+        edges: list of graph edges.
+
+    Raises:
+        ValueError: if the canvas dimensions are invalid or if
+            edge coordinates are invalid.
+
+    Returns:
+        ASCII representation
+
+    Example:
+        ```python
+        from coco_core.runnables.graph_ascii import draw_ascii
+
+        vertices = {1: "1", 2: "2", 3: "3", 4: "4"}
+        edges = [
+            (source, target, None, None)
+            for source, target in [(1, 2), (2, 3), (2, 4), (1, 4)]
+        ]
+
+
+        print(draw_ascii(vertices, edges))
+        ```
+
+        ```txt
+
+                 +---+
+                 | 1 |
+                 +---+
+                 *    *
+                *     *
+               *       *
+            +---+       *
+            | 2 |       *
+            +---+**     *
+              *    **   *
+              *      ** *
+              *        **
+            +---+     +---+
+            | 3 |     | 4 |
+            +---+     +---+
+        ```
+    """
+    # NOTE: coordinates might me negative, so we need to shift
+    # everything to the positive plane before we actually draw it.
+    xlist: list[float] = []
+    ylist: list[float] = []
+
+    sug = _build_sugiyama_layout(vertices, edges)
+
+    for vertex in sug.g.sV:
+        # NOTE: moving boxes w/2 to the left
+        xlist.extend(
+            (
+                vertex.view.xy[0] - vertex.view.w / 2.0,
+                vertex.view.xy[0] + vertex.view.w / 2.0,
+            )
+        )
+        ylist.extend((vertex.view.xy[1], vertex.view.xy[1] + vertex.view.h))
+
+    for edge in sug.g.sE:
+        for x, y in edge.view.pts:
+            xlist.append(x)
+            ylist.append(y)
+
+    minx = min(xlist)
+    miny = min(ylist)
+    maxx = max(xlist)
+    maxy = max(ylist)
+
+    canvas_cols = math.ceil(math.ceil(maxx) - math.floor(minx)) + 1
+    canvas_lines = round(maxy - miny)
+
+    canvas = AsciiCanvas(canvas_cols, canvas_lines)
+
+    # NOTE: first draw edges so that node boxes could overwrite them
+    for edge in sug.g.sE:
+        if len(edge.view.pts) <= 1:
+            msg = "Not enough points to draw an edge"
+            raise ValueError(msg)
+        for index in range(1, len(edge.view.pts)):
+            start = edge.view.pts[index - 1]
+            end = edge.view.pts[index]
+
+            start_x = round(start[0] - minx)
+            start_y = round(start[1] - miny)
+            end_x = round(end[0] - minx)
+            end_y = round(end[1] - miny)
+
+            if start_x < 0 or start_y < 0 or end_x < 0 or end_y < 0:
+                msg = (
+                    "Invalid edge coordinates: "
+                    f"start_x={start_x}, "
+                    f"start_y={start_y}, "
+                    f"end_x={end_x}, "
+                    f"end_y={end_y}"
+                )
+                raise ValueError(msg)
+
+            canvas.line(start_x, start_y, end_x, end_y, "." if edge.data else "*")
+
+    for vertex in sug.g.sV:
+        # NOTE: moving boxes w/2 to the left
+        x = vertex.view.xy[0] - vertex.view.w / 2.0
+        y = vertex.view.xy[1]
+
+        canvas.box(
+            round(x - minx),
+            round(y - miny),
+            vertex.view.w,
+            vertex.view.h,
+        )
+
+        canvas.text(round(x - minx) + 1, round(y - miny) + 1, vertex.data)
+
+    return canvas.draw()