pythagoras_tree.py
1 from PyQt5.QtCore import * 2 from PyQt5.QtGui import * 3 from PyQt5.QtWidgets import * 4 5 from vector import * 6 7 ANGLE = 30 8 9 10 class Square: 11 def __init__(self, p1: Point, p2: Point): 12 self.p1 = p1 13 self.p2 = p2 14 vector = (p2 - p1).to_vector() 15 vector.angle -= 90 16 self.p3 = p2 + vector.to_point() 17 self.p4 = p1 + vector.to_point() 18 19 self.triangle = None 20 21 def __iter__(self): 22 yield self.p1, self.p2, self.p3, self.p4 23 if self.triangle: 24 yield from self.triangle 25 26 def iterate(self): 27 if self.triangle is None: 28 self.triangle = Triangle(self.p4, self.p3) 29 else: 30 self.triangle.iterate() 31 32 33 class Triangle: 34 def __init__(self, p1: Point, p2: Point): 35 self.p1 = p1 36 self.p2 = p2 37 vector = (p2 - p1).to_vector() 38 vector.angle -= ANGLE 39 vector.distance *= math.cos(deg2rad(ANGLE)) 40 self.p3 = p1 + vector.to_point() 41 42 self.square1 = self.square2 = None 43 44 def __iter__(self): 45 yield self.p1, self.p2, self.p3 46 if self.square1 is not None: 47 yield from self.square1 48 yield from self.square2 49 50 def iterate(self): 51 if self.square1 is None: 52 self.square1 = Square(self.p1, self.p3) 53 self.square2 = Square(self.p3, self.p2) 54 else: 55 self.square1.iterate() 56 self.square2.iterate() 57 58 59 class PyTree(QWidget): 60 def __init__(self): 61 super().__init__() 62 63 self.setWindowTitle("Pythagoras Tree - Iteration 0") 64 self.setFixedSize(900, 600) 65 66 self.tree = Square( 67 Point(self.width() / 2 - 50, self.height() - 20), 68 Point(self.width() / 2 + 50, self.height() - 20) 69 ) 70 self.iterations = [list(self.tree)] 71 self.current_iteration = 0 72 73 self.show() 74 75 def keyPressEvent(self, e: QKeyEvent): 76 if e.key() == Qt.Key_Q: 77 self.close() 78 elif e.key() == Qt.Key_Right: 79 self.current_iteration += 1 80 if len(self.iterations) <= self.current_iteration: 81 self.tree.iterate() 82 self.iterations.append(list(self.tree)) 83 self.setWindowTitle(f"Pythagoras Tree - Iteration {self.current_iteration}") 84 self.repaint() 85 elif e.key() == Qt.Key_Left: 86 if self.current_iteration > 0: 87 self.current_iteration -= 1 88 self.setWindowTitle(f"Pythagoras Tree - Iteration {self.current_iteration}") 89 self.repaint() 90 91 def paintEvent(self, _: QPaintEvent): 92 qp = QPainter(self) 93 qp.setPen(Qt.white) 94 qp.setBrush(Qt.white) 95 qp.drawRect(self.rect()) 96 97 qp.setPen(Qt.black) 98 99 for polygon in self.iterations[self.current_iteration]: 100 qp.drawPolygon(*(QPoint(*p) for p in polygon)) 101 102 103 if __name__ == '__main__': 104 app = QApplication([]) 105 tree = PyTree() 106 app.exec_()