BinaryTreeDiagram

src.drawpyo.diagram_types.binary_tree.BinaryTreeDiagram

Bases: TreeDiagram

Simplifies TreeDiagram for binary-tree convenience.

Source code in src/drawpyo/diagram_types/binary_tree.py

class BinaryTreeDiagram(TreeDiagram):
    """Simplifies TreeDiagram for binary-tree convenience."""

    DEFAULT_LEVEL_SPACING = 80
    DEFAULT_ITEM_SPACING = 20
    DEFAULT_GROUP_SPACING = 30
    DEFAULT_LINK_STYLE = "straight"

    def __init__(self, **kwargs) -> None:
        kwargs.setdefault("level_spacing", self.DEFAULT_LEVEL_SPACING)
        kwargs.setdefault("item_spacing", self.DEFAULT_ITEM_SPACING)
        kwargs.setdefault("group_spacing", self.DEFAULT_GROUP_SPACING)
        kwargs.setdefault("link_style", self.DEFAULT_LINK_STYLE)
        super().__init__(**kwargs)

    def _attach(
        self, parent: BinaryNodeObject, child: BinaryNodeObject, side: str
    ) -> None:
        if not isinstance(parent, BinaryNodeObject) or not isinstance(
            child, BinaryNodeObject
        ):
            raise TypeError("parent and child must be BinaryNodeObject instances")

        if parent.tree is not self:
            parent.tree = self
        child.tree = self

        setattr(parent, side, child)

    def add_left(self, parent: BinaryNodeObject, child: BinaryNodeObject) -> None:
        self._attach(parent, child, "left")

    def add_right(self, parent: BinaryNodeObject, child: BinaryNodeObject) -> None:
        self._attach(parent, child, "right")

    @classmethod
    def from_dict(
        cls,
        data: dict,
        *,
        colors: list = None,
        coloring: str = "depth",
        **kwargs,
    ) -> "BinaryTreeDiagram":
        """
        Build a BinaryTreeDiagram from nested dict/list structures.
        data: Nested dict/list structure representing the tree.
        colors: List of ColorSchemes, StandardColors, or color hex strings to use for coloring nodes. Default: None
        coloring: str - "depth" | "hash" | "type" | "directional" - Method to match colors to nodes. Default: "depth"
            1. "depth" - Color nodes based on their depth in the tree.
            2. "hash" - Color nodes based on a hash of their value.
            3. "type" - Color nodes based on their type (category, list_item, leaf).
            4. "directional" - Color nodes based on their direction (left, right).

            Note: In colors Left Nodes are coloured by 0th index and Right Nodes are coloured by 1st index in the list of colors.
        """

        if coloring not in {"depth", "hash", "type", "directional"}:
            raise ValueError(f"Invalid coloring mode: {coloring}")

        if colors is not None and not isinstance(colors, list):
            raise TypeError("colors must be a list or None")

        colors = colors or None
        TYPE_INDEX = {"category": 0, "list_item": 1, "leaf": 2}

        # -------------------------
        # Validation
        # -------------------------

        def validate(tree_node: Dict, *, is_root=False):
            if tree_node is None or isinstance(tree_node, (str, int, float)):
                return

            if isinstance(tree_node, (list, tuple)):
                if len(tree_node) > 2:
                    raise TypeError("List node can have at most two children")
                for x in tree_node:
                    validate(x)
                return

            if isinstance(tree_node, dict):
                if is_root and len(tree_node) != 1:
                    raise TypeError("Root dict must contain exactly one key")

                if not is_root and not (1 <= len(tree_node) <= 2):
                    raise TypeError("Dict node must have 1 or 2 children")

                for node, children in tree_node.items():
                    if not isinstance(node, (str, int, float)):
                        raise TypeError(f"Invalid dict key type: {type(node)}")

                    validate(children)
                return

            raise TypeError(f"Unsupported tree tree_node type: {type(tree_node)}")

        if not isinstance(data, dict):
            raise TypeError("Top-level tree must be a dict")

        # Checks if the provided dict data is valid for a binary tree construction
        validate(data, is_root=True)

        # -------------------------
        # Helpers
        # -------------------------

        diagram = cls(**kwargs)

        def choose_color(
            value: str, node_type: str, depth: int, side: Optional[object] = None
        ):
            if not colors:
                return None

            n = len(colors)

            if coloring == "depth":
                idx = depth % n
            elif coloring == "hash":
                h = int(hashlib.md5(value.encode()).hexdigest(), 16)
                idx = h % n
            elif coloring == "directional":
                # side can be 'left'/'right' or a boolean where True==left
                if n != 2:
                    raise ValueError(
                        "colors list must be of length atleast 2 for directional coloring"
                    )

                if side is None:
                    return None

                if isinstance(side, bool):
                    is_left = side
                else:
                    is_left = str(side).lower() == "left"

                idx = (0 if is_left else 1) % n

            else:  # type
                idx = TYPE_INDEX[node_type] % n

            return colors[idx]

        def create_node(value: str, parent, color):
            if color is None:
                return BinaryNodeObject(tree=diagram, value=value, tree_parent=parent)

            if isinstance(color, drawpyo.ColorScheme):
                return BinaryNodeObject(
                    tree=diagram, value=value, tree_parent=parent, color_scheme=color
                )

            return BinaryNodeObject(
                tree=diagram, value=value, tree_parent=parent, fillColor=color
            )

        # -------------------------
        # Build
        # -------------------------

        def build(parent: BinaryNodeObject, item: Any, depth: int):
            if item is None:
                return

            # Leaf
            if isinstance(item, (str, int, float)):
                value = str(item)
                # leaf nodes in this branch are always attached as left
                node = create_node(
                    value,
                    parent,
                    choose_color(value, "leaf", depth, side="left"),
                )
                diagram.add_left(parent, node)
                return

            # Dict (named children)
            if isinstance(item, dict):
                for index, (node, children) in enumerate(item.items()):
                    name = str(node)
                    side = "left" if index == 0 else "right"
                    node = create_node(
                        name,
                        parent,
                        choose_color(name, "category", depth, side=side),
                    )

                    if index == 0:
                        diagram.add_left(parent, node)
                    else:
                        diagram.add_right(parent, node)

                    build(node, children, depth + 1)
                return

            # List / Tuple (positional children)
            for index, elem in enumerate(item):
                if elem is None:
                    continue

                if isinstance(elem, (str, int, float)):
                    name = str(elem)
                    side = "left" if index == 0 else "right"
                    node = create_node(
                        name,
                        parent,
                        choose_color(name, "leaf", depth + 1, side=side),
                    )

                elif isinstance(elem, dict) and len(elem) == 1:
                    node, children = next(iter(elem.items()))
                    name = str(node)
                    side = "left" if index == 0 else "right"
                    node = create_node(
                        name,
                        parent,
                        choose_color(name, "category", depth + 1, side=side),
                    )
                    build(node, children, depth + 1)
                else:
                    raise TypeError(
                        "List elements must be primitive or single-key dict"
                    )

                if index == 0:
                    diagram.add_left(parent, node)
                else:
                    diagram.add_right(parent, node)

        # -------------------------
        # Root
        # -------------------------

        root_key, root_value = next(iter(data.items()))
        root_name = str(root_key)

        root = create_node(
            root_name,
            None,
            choose_color(root_name, "category", 0, None),
        )

        build(root, root_value, depth=1)
        diagram.auto_layout()
        return diagram

from_dict(data, *, colors=None, coloring='depth', **kwargs) classmethod

Build a BinaryTreeDiagram from nested dict/list structures. data: Nested dict/list structure representing the tree. colors: List of ColorSchemes, StandardColors, or color hex strings to use for coloring nodes. Default: None coloring: str - "depth" | "hash" | "type" | "directional" - Method to match colors to nodes. Default: "depth" 1. "depth" - Color nodes based on their depth in the tree. 2. "hash" - Color nodes based on a hash of their value. 3. "type" - Color nodes based on their type (category, list_item, leaf). 4. "directional" - Color nodes based on their direction (left, right).

Note: In colors Left Nodes are coloured by 0th index and Right Nodes are coloured by 1st index in the list of colors.

Source code in src/drawpyo/diagram_types/binary_tree.py

@classmethod
def from_dict(
    cls,
    data: dict,
    *,
    colors: list = None,
    coloring: str = "depth",
    **kwargs,
) -> "BinaryTreeDiagram":
    """
    Build a BinaryTreeDiagram from nested dict/list structures.
    data: Nested dict/list structure representing the tree.
    colors: List of ColorSchemes, StandardColors, or color hex strings to use for coloring nodes. Default: None
    coloring: str - "depth" | "hash" | "type" | "directional" - Method to match colors to nodes. Default: "depth"
        1. "depth" - Color nodes based on their depth in the tree.
        2. "hash" - Color nodes based on a hash of their value.
        3. "type" - Color nodes based on their type (category, list_item, leaf).
        4. "directional" - Color nodes based on their direction (left, right).

        Note: In colors Left Nodes are coloured by 0th index and Right Nodes are coloured by 1st index in the list of colors.
    """

    if coloring not in {"depth", "hash", "type", "directional"}:
        raise ValueError(f"Invalid coloring mode: {coloring}")

    if colors is not None and not isinstance(colors, list):
        raise TypeError("colors must be a list or None")

    colors = colors or None
    TYPE_INDEX = {"category": 0, "list_item": 1, "leaf": 2}

    # -------------------------
    # Validation
    # -------------------------

    def validate(tree_node: Dict, *, is_root=False):
        if tree_node is None or isinstance(tree_node, (str, int, float)):
            return

        if isinstance(tree_node, (list, tuple)):
            if len(tree_node) > 2:
                raise TypeError("List node can have at most two children")
            for x in tree_node:
                validate(x)
            return

        if isinstance(tree_node, dict):
            if is_root and len(tree_node) != 1:
                raise TypeError("Root dict must contain exactly one key")

            if not is_root and not (1 <= len(tree_node) <= 2):
                raise TypeError("Dict node must have 1 or 2 children")

            for node, children in tree_node.items():
                if not isinstance(node, (str, int, float)):
                    raise TypeError(f"Invalid dict key type: {type(node)}")

                validate(children)
            return

        raise TypeError(f"Unsupported tree tree_node type: {type(tree_node)}")

    if not isinstance(data, dict):
        raise TypeError("Top-level tree must be a dict")

    # Checks if the provided dict data is valid for a binary tree construction
    validate(data, is_root=True)

    # -------------------------
    # Helpers
    # -------------------------

    diagram = cls(**kwargs)

    def choose_color(
        value: str, node_type: str, depth: int, side: Optional[object] = None
    ):
        if not colors:
            return None

        n = len(colors)

        if coloring == "depth":
            idx = depth % n
        elif coloring == "hash":
            h = int(hashlib.md5(value.encode()).hexdigest(), 16)
            idx = h % n
        elif coloring == "directional":
            # side can be 'left'/'right' or a boolean where True==left
            if n != 2:
                raise ValueError(
                    "colors list must be of length atleast 2 for directional coloring"
                )

            if side is None:
                return None

            if isinstance(side, bool):
                is_left = side
            else:
                is_left = str(side).lower() == "left"

            idx = (0 if is_left else 1) % n

        else:  # type
            idx = TYPE_INDEX[node_type] % n

        return colors[idx]

    def create_node(value: str, parent, color):
        if color is None:
            return BinaryNodeObject(tree=diagram, value=value, tree_parent=parent)

        if isinstance(color, drawpyo.ColorScheme):
            return BinaryNodeObject(
                tree=diagram, value=value, tree_parent=parent, color_scheme=color
            )

        return BinaryNodeObject(
            tree=diagram, value=value, tree_parent=parent, fillColor=color
        )

    # -------------------------
    # Build
    # -------------------------

    def build(parent: BinaryNodeObject, item: Any, depth: int):
        if item is None:
            return

        # Leaf
        if isinstance(item, (str, int, float)):
            value = str(item)
            # leaf nodes in this branch are always attached as left
            node = create_node(
                value,
                parent,
                choose_color(value, "leaf", depth, side="left"),
            )
            diagram.add_left(parent, node)
            return

        # Dict (named children)
        if isinstance(item, dict):
            for index, (node, children) in enumerate(item.items()):
                name = str(node)
                side = "left" if index == 0 else "right"
                node = create_node(
                    name,
                    parent,
                    choose_color(name, "category", depth, side=side),
                )

                if index == 0:
                    diagram.add_left(parent, node)
                else:
                    diagram.add_right(parent, node)

                build(node, children, depth + 1)
            return

        # List / Tuple (positional children)
        for index, elem in enumerate(item):
            if elem is None:
                continue

            if isinstance(elem, (str, int, float)):
                name = str(elem)
                side = "left" if index == 0 else "right"
                node = create_node(
                    name,
                    parent,
                    choose_color(name, "leaf", depth + 1, side=side),
                )

            elif isinstance(elem, dict) and len(elem) == 1:
                node, children = next(iter(elem.items()))
                name = str(node)
                side = "left" if index == 0 else "right"
                node = create_node(
                    name,
                    parent,
                    choose_color(name, "category", depth + 1, side=side),
                )
                build(node, children, depth + 1)
            else:
                raise TypeError(
                    "List elements must be primitive or single-key dict"
                )

            if index == 0:
                diagram.add_left(parent, node)
            else:
                diagram.add_right(parent, node)

    # -------------------------
    # Root
    # -------------------------

    root_key, root_value = next(iter(data.items()))
    root_name = str(root_key)

    root = create_node(
        root_name,
        None,
        choose_color(root_name, "category", 0, None),
    )

    build(root, root_value, depth=1)
    diagram.auto_layout()
    return diagram