Color Label Index for Edges instead of Labels

[ryanpeach]

Hey I just wanted to share something I did quickly today:

I created code that overwrote the Diagram object like so:

from typing import Dict, List, Set
from diagrams import Cluster as _Cluster
from diagrams import Diagram as _Diagram
from diagrams import Edge as _Edge
import colorsys
import random

LABELS_TO_COLORS: Dict[str, str] = {}

def generate_list_of_color_options(s=100.0, l=50.0) -> Generator[str, None, None]:
    """
    Algorithmically generate an infinite list of numbers which are evenly spaced apart from each other.
    Then, convert each number to a color using the HSL color space.
    Invented myself and not tested so don't quote me on this, but the way it works is:
      Imagine a color wheel with 360 degrees and the angle in degrees being `h`.
      Start off subdividing the color wheel into 3 equal parts.
      Iterate over those parts yielding the color at each angle.
      Then when you iterate over the full circle, subdivide the parts into 2 times as many parts, and iterate over those.
      In this case some colors will duplicate, eliminate those. This could be done mathematically but I'm lazy, so I did it with a set.
      Because I used a set you need to round the numbers to integers, or else floating point errors will cause near duplicates to not be eliminated.
      Stop after 360 colors have been generated, because that's the maximum number of colors you can have in a 360 degree color wheel under just the number h.
    You can modify s and l to change the saturation and lightness of the colors.
    """
    h = 0
    increment: float = 1.0 / 3.0 * 360
    seen_hues: Set[int] = set()
    while True:
        h = round(h + increment)
        if h > 360:
            h -= 360
            increment /= 2.0
        if h not in seen_hues:
            seen_hues.add(h)
            rgb = colorsys.hls_to_rgb(h / 360, l / 100, s / 100)
            yield "#%02x%02x%02x" % (
                round(rgb[0] * 255),
                round(rgb[1] * 255),
                round(rgb[2] * 255),
            )
        if len(seen_hues) == 360 or increment < 1:
            raise Exception("Ran out of colors")

LIST_OF_COLOR_OPTIONS = generate_list_of_color_options()

# Overwrite the Edge name
def Edge(**kwargs) -> _Edge:
    """Overwrite the Edge function to have one name for label, and to automatically provide other arguments to the Edge function."""
    label = kwargs.get("label", None) or kwargs.get("xlabel", None)
    if label in LABELS_TO_COLORS:
        color = LABELS_TO_COLORS[label]
    else:
        color = next(LIST_OF_COLOR_OPTIONS)
        LABELS_TO_COLORS[label] = color
    return _Edge(color=color)

class Diagram:

    def __init__(self, txt, **kwargs):
        self.diagram = _Diagram(txt, show=False, direction="TB")

    def __enter__(self):
        return self.diagram.__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        overwrite_edge_label_index(self.diagram)
        return self.diagram.__exit__(exc_type, exc_val, exc_tb)

# Write an index of colors to edge label
def overwrite_edge_label_index(diagram: _Diagram):
    # Cell text is going to be like
    # Empty Space | Label
    # Empty Space | Label
    # Empty Space | Label
    cellText = []
    for label, color in LABELS_TO_COLORS.items():
        cellText.append([" ", label])
    # Cell colors is going to be like
    # Color | White
    # Color | White
    # Color | White
    cellColors = []
    for label, color in LABELS_TO_COLORS.items():
        cellColors.append([color, "#FFFFFF"])

    # Create a table in HTML to display each col/row label and its color
    html = "<<TABLE>"
    for row in range(len(cellText)):
        html += "<TR>\n"
        for col in range(len(cellText[row])):
            html += f"<TD BGCOLOR=\"{cellColors[row][col]}\">{cellText[row][col]}</TD>\n"
        html += "</TR>\n"
    html += "</TABLE>>"

    diagram.dot.node(
        name="legend_fy23gdbqqbs",
        label=html,
        shape="plaintext",
        height=str(len(cellText) * 0.3),
    )

Please ignore the horrible chicken scratch code XD.

Anyway, this has the cool effect of eliminating edge labels in your diagram, and instead picking random colors from this list and assigning them to the edges, then adding an HTML Table to the graphviz which tells the user what code is associated with which labels.

Here is an example:

from my_library import Cluster, Diagram, Edge
from diagrams.k8s.compute import Deploy, Job
from diagrams.k8s.controlplane import API
from diagrams.onprem.inmemory import Redis

with Diagram("Example"):
    deploy = Deploy("Some Deployment")
    job = Job("Some Job")
    redis = Redis("Redis")
    api = API("K8s API")
    job >> Edge(label="writes") >> redis
    deploy >> Edge(label="reads") >> redis
    deploy >> Edge(label="posts") >> api
    api >> Edge(label="creates") >> job

I know this isn't an "issue" but maybe if this became a feature request we could automate this sort of thing and make it a Diagram option. I may have time to contribute but I'm unsure. Maybe people can upvote this to see if its in demand. Otherwise it'll just live on the internet as copy pasta for a future person wanting to do this as well.

[laxdog]

Thanks @ryanpeach , this is exactly what I needed, but couldn't figure out how to do.