ScoDoc/tools/create_api_map.py

import xml.etree.ElementTree as ET


class COLORS:
    BLUE = "rgb(114,159,207)"
    GREEN = "rgb(165,214,165)"
    PINK = "rgb(230,156,190)"
    GREY = "rgb(224,224,224)"
    ORANGE = "rgb(253,191,111)"


class Token:
    def __init__(self, name, method="GET", query=None, leaf=False):
        self.children: list["Token"] = []
        self.name: str = name
        self.method: str = method
        self.query: dict[str, str] = query or {}
        self.force_leaf: bool = leaf

    def add_child(self, child):
        self.children.append(child)

    def find_child(self, name):
        for child in self.children:
            if child.name == name:
                return child
        return None

    def __repr__(self, level=0):
        ret = "\t" * level + f"({self.name})\n"
        for child in self.children:
            ret += child.__repr__(level + 1)
        return ret

    def is_leaf(self):
        return len(self.children) == 0

    def get_height(self, y_step):
        # Calculer la hauteur totale des enfants
        children_height = sum(child.get_height(y_step) for child in self.children)

        # Calculer la hauteur des éléments de la query
        query_height = len(self.query) * y_step

        # La hauteur totale est la somme de la hauteur des enfants et des éléments de la query
        return children_height + query_height + y_step

    def to_svg_group(
        self,
        x_offset=0,
        y_offset=0,
        x_step=150,
        y_step=50,
        parent_coords=None,
        parent_children_nb=0,
    ):
        group = ET.Element("g")
        color = COLORS.BLUE
        if self.force_leaf or self.is_leaf():
            if self.method == "GET":
                color = COLORS.GREEN
            elif self.method == "POST":
                color = COLORS.PINK

        element = _create_svg_element(self.name, color)
        element.set("transform", f"translate({x_offset}, {y_offset})")
        current_start_coords, current_end_coords = _get_anchor_coords(
            element, x_offset, y_offset
        )
        group.append(element)

        # Add an arrow from the parent element to the current element
        if parent_coords and parent_children_nb > 1:
            arrow = _create_arrow(parent_coords, current_start_coords)
            group.append(arrow)

        if not self.is_leaf():
            slash_group = _create_svg_element("/", COLORS.GREY)
            slash_group.set(
                "transform",
                f"translate({x_offset + _get_element_width(element)}, {y_offset})",
            )
            group.append(slash_group)
        if self.is_leaf() and self.query:
            slash_group = _create_svg_element("?", COLORS.GREY)
            slash_group.set(
                "transform",
                f"translate({x_offset + _get_element_width(element)}, {y_offset})",
            )
            group.append(slash_group)
        current_end_coords = _get_anchor_coords(group, 0, 0)[1]

        query_y_offset = y_offset
        ampersand_start_coords = None
        ampersand_end_coords = None
        for key, value in self.query.items():
            # <param>=<value>
            translate_x = x_offset + x_step

            # <param>
            param_el = _create_svg_element(key, COLORS.BLUE)
            param_el.set(
                "transform",
                f"translate({translate_x}, {query_y_offset})",
            )
            group.append(param_el)

            # add Arrow from "query" to element
            coords = (
                current_end_coords,
                _get_anchor_coords(param_el, translate_x, query_y_offset)[0],
            )
            group.append(_create_arrow(*coords))

            # =
            equal_el = _create_svg_element("=", COLORS.GREY)
            translate_x = x_offset + x_step + _get_element_width(param_el)
            equal_el.set(
                "transform",
                f"translate({translate_x}, {query_y_offset})",
            )
            group.append(equal_el)

            # <value>
            value_el = _create_svg_element(value, COLORS.GREEN)
            translate_x = (
                x_offset
                + x_step
                + sum(_get_element_width(el) for el in [param_el, equal_el])
            )
            value_el.set(
                "transform",
                f"translate({translate_x}, {query_y_offset})",
            )
            group.append(value_el)
            if len(self.query) == 1:
                continue
            ampersand_group = _create_svg_element("&", "rgb(224,224,224)")
            translate_x = (
                x_offset
                + x_step
                + sum(_get_element_width(el) for el in [param_el, equal_el, value_el])
            )
            ampersand_group.set(
                "transform",
                f"translate({translate_x}, {query_y_offset})",
            )
            group.append(ampersand_group)

            # Track the start and end coordinates of the ampersands
            if ampersand_start_coords is None:
                ampersand_start_coords = _get_anchor_coords(
                    ampersand_group, translate_x, query_y_offset
                )[1]
            ampersand_end_coords = _get_anchor_coords(
                ampersand_group, translate_x, query_y_offset
            )[1]
            # Draw line connecting all ampersands
            if ampersand_start_coords and ampersand_end_coords and len(self.query) > 1:
                line = _create_line(ampersand_start_coords, ampersand_end_coords)
                group.append(line)

            query_y_offset += y_step

        y_offset = query_y_offset

        current_y_offset = y_offset

        for child in self.children:
            rel_x_offset = x_offset + _get_group_width(group)
            if len(self.children) > 1:
                rel_x_offset += x_step
            child_group = child.to_svg_group(
                rel_x_offset,
                current_y_offset,
                x_step,
                y_step,
                parent_coords=current_end_coords,
                parent_children_nb=len(self.children),
            )
            group.append(child_group)
            current_y_offset += child.get_height(y_step)

        return group


def _create_line(start_coords, end_coords):
    start_x, start_y = start_coords
    end_x, end_y = end_coords

    path_data = f"M {start_x},{start_y} L {start_x + 20},{start_y} L {start_x + 20},{end_y} L {end_x},{end_y}"

    path = ET.Element(
        "path",
        {
            "d": path_data,
            "style": "stroke:black;stroke-width:2;fill:none",
        },
    )
    return path


def _create_svg_element(text, color="rgb(230,156,190)"):
    # Dimensions and styling
    padding = 5
    font_size = 16
    rect_height = 30
    rect_x = 10
    rect_y = 20

    # Estimate the text width
    text_width = (
        len(text) * font_size * 0.6
    )  # Estimate based on average character width
    rect_width = text_width + padding * 2

    # Create the SVG group element
    group = ET.Element("g")

    # Create the rectangle
    ET.SubElement(
        group,
        "rect",
        {
            "x": str(rect_x),
            "y": str(rect_y),
            "width": str(rect_width),
            "height": str(rect_height),
            "style": f"fill:{color};stroke:black;stroke-width:2;fill-opacity:1;stroke-opacity:1",
        },
    )

    # Create the text element
    text_element = ET.SubElement(
        group,
        "text",
        {
            "x": str(rect_x + padding),
            "y": str(
                rect_y + rect_height / 2 + font_size / 2.5
            ),  # Adjust to vertically center the text
            "font-family": "Courier New, monospace",
            "font-size": str(font_size),
            "fill": "black",
            "style": "white-space: pre;",
        },
    )
    text_element.text = text

    return group


def _get_anchor_coords(element, x_offset, y_offset):
    bbox = _get_bbox(element, x_offset, y_offset)
    startX = bbox["x_min"]
    endX = bbox["x_max"]
    y = bbox["y_min"] + (bbox["y_max"] - bbox["y_min"]) / 2
    return (startX, y), (endX, y)


def _create_arrow(start_coords, end_coords):
    start_x, start_y = start_coords
    end_x, end_y = end_coords
    mid_x = (start_x + end_x) / 2

    path_data = (
        f"M {start_x},{start_y} L {mid_x},{start_y} L {mid_x},{end_y} L {end_x},{end_y}"
    )

    path = ET.Element(
        "path",
        {
            "d": path_data,
            "style": "stroke:black;stroke-width:2;fill:none",
            "marker-end": "url(#arrowhead)",
        },
    )
    return path


def _get_element_width(element):
    rect = element.find("rect")
    if rect is not None:
        return float(rect.get("width", 0))
    return 0


def _get_group_width(group):
    return sum(_get_element_width(child) for child in group)


def gen_api_map(app):
    api_map = Token("")
    for rule in app.url_map.iter_rules():
        # On ne garde que les routes de l'API / APIWEB
        if not rule.endpoint.lower().startswith("api"):
            continue

        segments = rule.rule.strip("/").split("/")
        current_token = api_map

        for i, segment in enumerate(segments):
            # Check if the segment already exists in the current level
            child = current_token.find_child(segment)
            if child is None:
                # If it's the last segment and it has query parameters
                if i == len(segments) - 1 and segment == "query":
                    child = Token(
                        segment,
                        leaf=True,
                    )
                    # TODO Parse QUERY doc
                else:
                    child = Token(segment)
                method: str = "POST" if "POST" in rule.methods else "GET"
                child.method = method
                current_token.add_child(child)
            current_token = child

        # Mark the last segment as a leaf node if it's not already marked
        if not current_token.is_leaf():
            current_token.force_leaf = True

    generate_svg(api_map.to_svg_group(), "/tmp/api_map.svg")


def _get_bbox(element, x_offset=0, y_offset=0):
    # Helper function to calculate the bounding box of an SVG element
    bbox = {
        "x_min": float("inf"),
        "y_min": float("inf"),
        "x_max": float("-inf"),
        "y_max": float("-inf"),
    }

    for child in element:
        transform = child.get("transform")
        child_x_offset = x_offset
        child_y_offset = y_offset

        if transform:
            translate = transform.replace("translate(", "").replace(")", "").split(",")
            if len(translate) == 2:
                child_x_offset += float(translate[0])
                child_y_offset += float(translate[1])

        if child.tag == "rect":
            x = child_x_offset + float(child.get("x", 0))
            y = child_y_offset + float(child.get("y", 0))
            width = float(child.get("width", 0))
            height = float(child.get("height", 0))
            bbox["x_min"] = min(bbox["x_min"], x)
            bbox["y_min"] = min(bbox["y_min"], y)
            bbox["x_max"] = max(bbox["x_max"], x + width)
            bbox["y_max"] = max(bbox["y_max"], y + height)

        if len(child):
            child_bbox = _get_bbox(child, child_x_offset, child_y_offset)
            bbox["x_min"] = min(bbox["x_min"], child_bbox["x_min"])
            bbox["y_min"] = min(bbox["y_min"], child_bbox["y_min"])
            bbox["x_max"] = max(bbox["x_max"], child_bbox["x_max"])
            bbox["y_max"] = max(bbox["y_max"], child_bbox["y_max"])

    return bbox


def generate_svg(element, fname):
    bbox = _get_bbox(element)
    width = bbox["x_max"] - bbox["x_min"] + 20  # Add some padding
    height = bbox["y_max"] - bbox["y_min"] + 20  # Add some padding

    svg = ET.Element(
        "svg",
        {
            "width": str(width),
            "height": str(height),
            "xmlns": "http://www.w3.org/2000/svg",
            "viewBox": f"{bbox['x_min'] - 10} {bbox['y_min'] - 10} {width} {height}",
        },
    )

    # Define the marker for the arrowhead
    defs = ET.SubElement(svg, "defs")
    marker = ET.SubElement(
        defs,
        "marker",
        {
            "id": "arrowhead",
            "markerWidth": "10",
            "markerHeight": "7",
            "refX": "10",
            "refY": "3.5",
            "orient": "auto",
        },
    )
    ET.SubElement(marker, "polygon", {"points": "0 0, 10 3.5, 0 7"})

    svg.append(element)

    tree = ET.ElementTree(svg)
    tree.write(fname, encoding="utf-8", xml_declaration=True)


if __name__ == "__main__":
    root = Token("api")
    child1 = Token("assiduites", leaf=True)
    child2 = Token("count")
    child22 = Token("all")
    child23 = Token(
        "query", query={"param1": "value1", "param2": "value2", "param3": "value3"}
    )
    child3 = Token("justificatifs", "POST")

    root.add_child(child1)
    child1.add_child(child2)
    child2.add_child(child22)
    child2.add_child(child23)
    root.add_child(child3)

    group_element = root.to_svg_group()

    generate_svg(group_element, "/tmp/api_map.svg")