Source code for langchain_core.runnables.graph_ascii

"""Draws DAG in ASCII.
Adapted from https://github.com/iterative/dvc/blob/main/dvc/dagascii.py"""

import math
import os
from typing import Any, Mapping, Sequence

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
[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): if dx == 0: y = y0 else: y = y0 + int(round((x - x0) * dy / float(dx))) self.point(x, y, char) elif y0 < y1: for y in range(y0, y1 + 1): if dy == 0: x = x0 else: x = x0 + int(round((y - y0) * dx / float(dy))) self.point(x, y, char) else: for y in range(y1, y0 + 1): if dy == 0: x = x0 else: x = 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: raise ImportError( "Install grandalf to draw graphs: `pip install grandalf`." ) 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. Xs = [] # noqa: N806 Ys = [] # noqa: N806 sug = _build_sugiyama_layout(vertices, edges) for vertex in sug.g.sV: # NOTE: moving boxes w/2 to the left Xs.append(vertex.view.xy[0] - vertex.view.w / 2.0) Xs.append(vertex.view.xy[0] + vertex.view.w / 2.0) Ys.append(vertex.view.xy[1]) Ys.append(vertex.view.xy[1] + vertex.view.h) for edge in sug.g.sE: for x, y in edge.view._pts: Xs.append(x) Ys.append(y) minx = min(Xs) miny = min(Ys) maxx = max(Xs) maxy = max(Ys) 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()