| 448 | |
| 449 | |
| 450 | class Edge(): |
| 451 | def __init__(self, node0, node1): |
| 452 | self.node0 = node0 |
| 453 | self.node1 = node1 |
| 454 | |
| 455 | # When we draw the edge, we know the calling function is definitely not a leaf... |
| 456 | # and the called function is definitely not a trunk |
| 457 | node0.is_leaf = False |
| 458 | node1.is_trunk = False |
| 459 | |
| 460 | def __repr__(self): |
| 461 | return f"<Edge {self.node0} -> {self.node1}" |
| 462 | |
| 463 | def __lt__(self, other): |
| 464 | if self.node0 == other.node0: |
| 465 | return self.node1 < other.node1 |
| 466 | return self.node0 < other.node0 |
| 467 | |
| 468 | def to_dot(self): |
| 469 | ''' |
| 470 | Returns string format for embedding in a dotfile. Example output: |
| 471 | node_uid_a -> node_uid_b [color='#aaa' penwidth='2'] |
| 472 | :rtype: str |
| 473 | ''' |
| 474 | ret = self.node0.uid + ' -> ' + self.node1.uid |
| 475 | source_color = int(self.node0.uid.split("_")[-1], 16) % len(EDGE_COLORS) |
| 476 | ret += f' [color="{EDGE_COLORS[source_color]}" penwidth="2"]' |
| 477 | return ret |
| 478 | |
| 479 | def to_dict(self): |
| 480 | """ |
| 481 | :rtype: dict |
| 482 | """ |
| 483 | return { |
| 484 | 'source': self.node0.uid, |
| 485 | 'target': self.node1.uid, |
| 486 | 'directed': True, |
| 487 | } |
| 488 | |
| 489 | |
| 490 | class Group(): |