simetri.extensions.tree

  1"""Create a tree structure and draw it."""
  2
  3from typing import Sequence, Any
  4
  5import simetri.graphics as sg
  6
  7
  8diamond = sg.Shape([(0, 5), (3, 0), (0, -5), (-3, 0)], closed=True)
  9diamond.fill_color = sg.black
 10diamond.stroke = False
 11star2 = sg.regular_star_polygon(5, 2, 7)
 12star2.set_attribs('fill_color', sg.red)
 13star2.set_attribs('stroke', False)
 14square = sg.Shape([(0, 5), (5, 5), (5, 0), (0, 0)], closed=True)
 15star = sg.regular_star_polygon(8, 3, 7)
 16star.set_attribs('fill_color', sg.red)
 17star.set_attribs('stroke', False)
 18star3 = sg.regular_star_polygon(8, 2, 5)
 19star3.set_attribs('fill_color', sg.blue)
 20star3.set_attribs('stroke', False)
 21circle = sg.Circle((0, 0), 1.5, fill_color=sg.white, stroke=False)
 22hexagon = sg.Batch([sg.reg_poly_shape((0, 0), 6, 4, fill_color=sg.teal,
 23                                                    stroke=False), circle])
 24
 25
 26def next_id():
 27    next_id.counter += 1
 28    return next_id.counter
 29
 30
 31next_id.counter = 0
 32
 33
 34class TreeNode:
 35    """TreeNode object representing a tree structure.
 36    Each node has a tag, an id, a list of children, and optional extra attributes.
 37
 38    Args:
 39        tag (str): The tag of the node.
 40        children (list): A list of child nodes. Defaults to an empty list.
 41        extra (str): Optional extra attribute for the node.
 42        **kwargs: Additional keyword arguments for the node. (font_size, font_color, bold)
 43
 44    Returns:
 45        None
 46    """
 47
 48    def __init__(
 49        self,
 50        tag: str = "",
 51        children: Sequence["TreeNode"] = None,
 52        extra: Any = None,
 53        font_size = 12,
 54        font_color = sg.black,
 55        bold = False
 56    ):
 57        self.tag = tag
 58        self.id = next_id()
 59        self.children = children if children is not None else []
 60        self.extra = extra
 61        self.font_size = font_size
 62        self.font_color = font_color
 63        self.bold = bold
 64
 65    def add_child(self, child):
 66        """Adds a child node to the current node.
 67        Args:
 68            child (TreeNode): The child node to add.
 69
 70        Returns:
 71            None
 72        """
 73        if child.id not in [c.id for c in self.children]:
 74            self.children.append(child)
 75
 76    def num_all_children(self):
 77        """Counts the number of children and grandchildren of the node.
 78        Args:
 79            None
 80
 81        Returns:
 82            int: The number of children and grandchildren.
 83        """
 84        return len(self.children) + sum(
 85            child.num_all_children() for child in self.children
 86        )
 87
 88    def depth(self):
 89        """Calculates the depth of the node in the tree.
 90        Args:
 91            None
 92
 93
 94        Returns:
 95            int: The depth of the node.
 96        """
 97        if not self.children:
 98            return 0
 99        n = self.num_all_children()
100        m = self.children[-1].num_all_children()
101        return n - m
102
103
104def make_tree(
105    node,
106    canvas: "Canvas" = None,
107    file_path: str = None,
108    overwrite: bool = False,
109    dx: float = 10,
110    dy: float = 18,
111    icons = None,
112    line1_color = sg.gray,
113    line1_width = 1.75,
114    line1_cap = sg.LineCap.ROUND,
115    line2_color = sg.gray,
116    line2_width = 1,
117    line2_cap = sg.LineCap.ROUND
118):
119    """Creates a tree structure and draws it on the canvas.
120    Args:
121        node: The root node of the tree.
122        canvas: The canvas to draw the tree structure.
123        file_path: The file path to save the tree structure.
124        overwrite: Whether to overwrite the existing file.
125        dx: The horizontal distance between nodes.
126        dy: The vertical distance between nodes.
127        icons: A list of icons to use for the nodes.
128        line1_color: The color of the first line.
129        line1_width: The width of the first line.
130        line1_cap: The cap style of the first line.
131        line2_color: The color of the second line.
132        line2_width: The width of the second line.
133        line2_cap: The cap style of the second line.
134
135    Returns:
136        None
137"""
138    count = 0
139    if icons is None:
140        icons = [star, diamond, star3, hexagon]
141
142    icon1, icon2, icon3, icon4 = icons
143
144    def print_tree(node, indent: int = 0, canvas: "Canvas" = None):
145        """Prints the tree structure of the node and its children.
146        Args:
147            node: The node to print.
148            indent: The indentation level for the current node.
149            canvas: The canvas to draw the tree structure.
150            file_path: The file path to save the tree structure.
151            overwrite: Whether to overwrite the existing file.
152
153        Returns:
154            None
155        """
156        nonlocal count
157        count += 1
158        x = indent * dx
159        y = -count * dy
160        x1 = x
161        y1 = y - node.depth() * dy
162        cx, cy = x, y
163        if node.depth() > 0:
164            canvas.line(
165                (x, y),
166                (x1, y1),
167                line_color=line1_color,
168                line_width=line1_width,
169                line_cap=line1_cap,
170            )
171
172
173        x2 = x1 + indent * dx
174        y2 = y
175        if indent > 0:
176            x -= dx
177        p1 = (x, y)
178        p2 = (x2, y2)
179        if not sg.close_points2(p1, p2):
180            canvas.line(
181                p1,
182                p2,
183                line_color=line2_color,
184                line_width=line2_width,
185                line_cap=line2_cap,
186            )
187        if node.depth() > 0 and indent > 0:
188            icon3.move_to((cx, cy))
189            canvas.draw(icon3)
190        font_size = node.font_size
191        font_color = node.font_color
192        bold = node.bold
193
194        if indent > 1:
195            #
196            icon4.move_to((x, y))
197            canvas.draw(icon4)
198        elif indent == 1:
199            icon2.move_to((x, y))
200            canvas.draw(icon2)
201        else:
202            icon1.move_to((x, y))
203            canvas.draw(icon1)
204        canvas.text(
205            node.tag,
206            (x2, y2),
207            font_family=sg.FontFamily.MONOSPACE,
208            font_size=font_size,
209            color=sg.red,
210            anchor=sg.Anchor.WEST,
211            fill=False,
212            bold=bold,
213            font_color=font_color,
214        )
215        for child in node.children:
216            print_tree(child, indent + 1, canvas=canvas)
217
218    print_tree(node, canvas=canvas)
219    canvas.save(file_path, overwrite=overwrite)
220
221
222# canvas = sg.Canvas()
223# root = TreeNode("{} Base", extra="root", font_color=sg.orange)
224# methods = TreeNode("Methods", font_color=sg.blue)
225# root.add_child(methods)
226# transforms = [
227#     "translate",
228#     "rotate",
229#     "mirror",
230#     "glide",
231#     "scale",
232#     "shear",
233#     "transform",
234# ]
235# args = [
236#     "(dx: float, dy: float)",
237#     "(angle: float, about: Point)",
238#     "(about: Line)",
239#     "(glide_line: Line, glide_dist: float)",
240#     "(scale_x: float, scale_y: float)",
241#     "(shear_x:float, shear_y:)",
242#     "(transform_matrix: ndarray)",
243# ]
244
245# for i, trans in enumerate(transforms):
246#     methods.add_child(TreeNode(f"{trans}{args[i][:-1]}, reps: int=0) -> Self"))
247
248# make_tree(
249#     root, canvas=canvas, file_path="c:/tmp/tree_generator4.pdf", overwrite=True
250# )