Drawing Arrows in PyQt5 QTextEdit Relative to Text Content

Question

I'm working on a PyQt5 application where I need to draw arrows in a QTextEdit widget. The arrows represent relationships between specific words in the text. I have implemented a custom QTextEdit class and overridden the paintEvent method to draw the arrows. The arrows are drawn based on start and end positions, representing the word indices in the text.

Here's my current implementation:

import sys
from PyQt5.QtWidgets import QApplication, QTextEdit, QMainWindow, QVBoxLayout, QWidget, QPushButton
from PyQt5.QtGui import QPainter, QPen, QTextCursor, QColor
from PyQt5.QtCore import Qt, QPoint

class CustomTextEdit(QTextEdit):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.arrows = []

    def addArrow(self, start_pos, end_pos, color):
        self.arrows.append((start_pos, end_pos, color))
        self.updateArrows()
        self.repaint()

    def deleteArrow(self, start_pos, end_pos, color):
        self.arrows.remove((start_pos, end_pos, color))
        self.repaint()

    def resizeEvent(self, event):
        super().resizeEvent(event)
        self.updateArrows()

    def updateArrows(self):
        self.updated_arrows = []
        for start_pos, end_pos, color in self.arrows:
            cursor = self.textCursor()
            cursor.movePosition(QTextCursor.Start)
            cursor.movePosition(QTextCursor.Right, QTextCursor.KeepAnchor, start_pos)
            start_point = self.cursorRect(cursor).center()

            cursor.movePosition(QTextCursor.Start)
            cursor.movePosition(QTextCursor.Right, QTextCursor.KeepAnchor, end_pos)
            end_point = self.cursorRect(cursor).center()

            self.updated_arrows.append((start_point, end_point, color))
        self.repaint()

    def paintEvent(self, event):
        super().paintEvent(event)
        if hasattr(self, 'updated_arrows'):
            painter = QPainter(self.viewport())
            painter.setRenderHint(QPainter.Antialiasing)

            for start_point, end_point, color in self.updated_arrows:
                pen = QPen(QColor(color), 2, Qt.SolidLine)
                painter.setPen(pen)

                dx = end_point.x() - start_point.x()
                dy = end_point.y() - start_point.y()
                length = (dx ** 2 + dy ** 2) ** 0.5

                if length > 0:
                    dx /= length
                    dy /= length

                    # Draw line
                    painter.drawLine(start_point, end_point)

                    # Draw arrow
                    arrow_size = 5
                    line_length = 10
                    arrow_p1 = end_point - QPoint(int(line_length * dx - arrow_size * dy),
                                                 int(line_length * dy + arrow_size * dx))
                    arrow_p2 = end_point - QPoint(int(line_length * dx + arrow_size * dy),
                                                 int(line_length * dy - arrow_size * dx))
                    painter.drawLine(arrow_p1, end_point)
                    painter.drawLine(arrow_p2, end_point)
                    print(arrow_p1.x(), end_point.x(), color)

    def resetTextEdit(self):
        self.clear()  # Clear the text content
        self.arrows = []  # Remove all arrows
        self.update()


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        central_widget = QWidget()
        layout = QVBoxLayout()
        self.text_edit = CustomTextEdit(self)
        self.text_edit.setText("This is a sample text to demonstrate arrows with different colors.")
        self.text_edit.addArrow(50, 100, 'green')
        layout.addWidget(self.text_edit)
        self.button = QPushButton('Add Arrow', self)
        self.button.clicked.connect(self.addArrow)
        layout.addWidget(self.button)
        central_widget.setLayout(layout)
        self.setCentralWidget(central_widget)

        self.setGeometry(100, 100, 800, 600)
        self.setWindowTitle('Arrows with Different Colors')
        self.show()

    def addArrow(self):
        self.text_edit.addArrow(10, 50, 'red')
        print("Arrow added")

if __name__ == '__main__':
    app = QApplication(sys.argv)
    mainWindow = MainWindow()
    sys.exit(app.exec_())

The arrows are drawn correctly, when adding them directly after initialising the CustomTextEdit. However, when trying to add an arrow by clicking the QPushButton the arrow is not drawn to the viewport.

How can I ensure that the arrows are drawn even when adding them via QPushButton?

Thanks in advance for your help!