"""Draws DAG in ASCII.
Adapted from https://github.com/iterative/dvc/blob/main/dvc/dagascii.py"""
import math
import os
from collections.abc import Mapping, Sequence
from typing import Any
from langchain_core.runnables.graph import Edge as LangEdge
[docs]
class VertexViewer:
"""Class to define vertex box boundaries that will be accounted for during
graph building by grandalf.
Args:
name (str): name of the vertex.
"""
HEIGHT = 3 # top and bottom box edges + text
"""Height of the box."""
[docs]
def __init__(self, name: str) -> None:
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
[docs]
class AsciiCanvas:
"""Class for drawing in ASCII.
Args:
cols (int): number of columns in the canvas. Should be > 1.
lines (int): number of lines in the canvas. Should be > 1.
"""
TIMEOUT = 10
[docs]
def __init__(self, cols: int, lines: int) -> None:
assert cols > 1
assert lines > 1
self.cols = cols
self.lines = lines
self.canvas = [[" "] * cols for line in range(lines)]
[docs]
def draw(self) -> str:
"""Draws ASCII canvas on the screen."""
lines = map("".join, self.canvas)
return os.linesep.join(lines)
[docs]
def point(self, x: int, y: int, char: str) -> None:
"""Create a point on ASCII canvas.
Args:
x (int): x coordinate. Should be >= 0 and < number of columns in
the canvas.
y (int): y coordinate. Should be >= 0 an < number of lines in the
canvas.
char (str): character to place in the specified point on the
canvas.
"""
assert len(char) == 1
assert x >= 0
assert x < self.cols
assert y >= 0
assert y < self.lines
self.canvas[y][x] = char
[docs]
def line(self, x0: int, y0: int, x1: int, y1: int, char: str) -> None:
"""Create a line on ASCII canvas.
Args:
x0 (int): x coordinate where the line should start.
y0 (int): y coordinate where the line should start.
x1 (int): x coordinate where the line should end.
y1 (int): y coordinate where the line should end.
char (str): 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 + int(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 + int(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 + int(round((y - y1) * dx / float(dy)))
self.point(x, y, char)
[docs]
def text(self, x: int, y: int, text: str) -> None:
"""Print a text on ASCII canvas.
Args:
x (int): x coordinate where the text should start.
y (int): y coordinate where the text should start.
text (str): string that should be printed.
"""
for i, char in enumerate(text):
self.point(x + i, y, char)
[docs]
def box(self, x0: int, y0: int, width: int, height: int) -> None:
"""Create a box on ASCII canvas.
Args:
x0 (int): x coordinate of the box corner.
y0 (int): y coordinate of the box corner.
width (int): box width.
height (int): box height.
"""
assert width > 1
assert height > 1
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, "+")
def _build_sugiyama_layout(
vertices: Mapping[str, str], edges: Sequence[LangEdge]
) -> Any:
try:
from grandalf.graphs import Edge, Graph, Vertex # type: ignore[import]
from grandalf.layouts import SugiyamaLayout # type: ignore[import]
from grandalf.routing import ( # type: ignore[import]
EdgeViewer,
route_with_lines,
)
except ImportError as exc:
msg = "Install grandalf to draw graphs: `pip install grandalf`."
raise ImportError(msg) from exc
#
# 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
[docs]
def draw_ascii(vertices: Mapping[str, str], edges: Sequence[LangEdge]) -> str:
"""Build a DAG and draw it in ASCII.
Args:
vertices (list): list of graph vertices.
edges (list): list of graph edges.
Returns:
str: ASCII representation
Example:
>>> vertices = [1, 2, 3, 4]
>>> edges = [(1, 2), (2, 3), (2, 4), (1, 4)]
>>> print(draw(vertices, edges))
+---+ +---+
| 3 | | 4 |
+---+ *+---+
* ** *
* ** *
* * *
+---+ *
| 2 | *
+---+ *
* *
* *
**
+---+
| 1 |
+---+
"""
# NOTE: coordinates might me negative, so we need to shift
# everything to the positive plane before we actually draw it.
xlist = []
ylist = []
sug = _build_sugiyama_layout(vertices, edges)
for vertex in sug.g.sV:
# NOTE: moving boxes w/2 to the left
xlist.append(vertex.view.xy[0] - vertex.view.w / 2.0)
xlist.append(vertex.view.xy[0] + vertex.view.w / 2.0)
ylist.append(vertex.view.xy[1])
ylist.append(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 = int(math.ceil(math.ceil(maxx) - math.floor(minx))) + 1
canvas_lines = int(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:
assert len(edge.view._pts) > 1
for index in range(1, len(edge.view._pts)):
start = edge.view._pts[index - 1]
end = edge.view._pts[index]
start_x = int(round(start[0] - minx))
start_y = int(round(start[1] - miny))
end_x = int(round(end[0] - minx))
end_y = int(round(end[1] - miny))
assert start_x >= 0
assert start_y >= 0
assert end_x >= 0
assert end_y >= 0
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(
int(round(x - minx)),
int(round(y - miny)),
vertex.view.w,
vertex.view.h,
)
canvas.text(int(round(x - minx)) + 1, int(round(y - miny)) + 1, vertex.data)
return canvas.draw()
