Trying to enable Word Wrap for html strings in QTableView using QTextDocument

I utilized the following HTML painter class from a previous post for my QTableView. The delegate uses a QTextdocument to display HTML inside a QTableview. The issue is that HTML strings will not word wrap.

Before using the delegate/HTML strings, word wrap works:

However when I use HTML Painter class/delegate it applies HTML tags but doesn't word wrap:

There was another post which addressed this, and I tried adding settings for QTexDocument for word wrap, which doesn't seem to be working. Here is the HTML painter class and the two lines I changed:

class HtmlPainter(QtWidgets.QStyledItemDelegate):
    def __init__(self, parent=None):
        QtWidgets.QStyledItemDelegate.__init__(self, parent)
    def paint(self, painter, option, index):
        if index.column() == 1: 
            text = index.model().data(index) #default role is display
            palette = QtWidgets.QApplication.palette()
            document = QtGui.QTextDocument()
            document.setDefaultFont(option.font)
            
            # MY ADDED CODE FOR WORD WRAP
            
            textOption=QtGui.QTextOption(document.defaultTextOption())
            textOption.setWrapMode(QtGui.QTextOption.WordWrap)
           
            # MY ADDED CODE FOR WORD WRAP

            # Set text (color depends on whether selected)
            if option.state & QtWidgets.QStyle.State_Selected:  
                displayString = "<font color={0}>{1}</font>".format(palette.highlightedText().color().name(), text) 
                document.setHtml(displayString)
            else:
                document.setHtml(text)
            #Set background color
            bgColor = palette.highlight().color() if (option.state & QtWidgets.QStyle.State_Selected)\
                     else palette.base().color()
            painter.save()
            painter.fillRect(option.rect, bgColor)
            painter.translate(option.rect.x(), option.rect.y()+5)  #If I add +5 it works
            document.drawContents(painter)
            painter.restore()
        else:
            QtWidgets.QStyledItemDelegate.paint(self, painter, option, index)    

I needed to add `document.setTextWidth(option.rect.width()) and that worked

document.setTextWidth(option.rect.width())
painter.save()
painter.fillRect(option.rect, bgColor)