72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
import numpy
|
|
import pandas as pd
|
|
from PyQt6 import QtCore
|
|
from PyQt6.QtCore import Qt, QModelIndex, QSortFilterProxyModel
|
|
|
|
|
|
def _get_str_dataframe(df: pd.DataFrame) -> pd.DataFrame:
|
|
"""Return a given dataframe with all values turned into strings.
|
|
|
|
When the data given to the PandasModel class contains non-strings,
|
|
an attached QTableView seems to respond rather slowly. This function
|
|
turns all data in the DataFrame into strings, yielding a better experience.
|
|
"""
|
|
return df.astype(str)
|
|
|
|
|
|
class PandasModel(QtCore.QAbstractTableModel):
|
|
def __init__(self, parent=None):
|
|
QtCore.QAbstractTableModel.__init__(self, parent)
|
|
|
|
self._data = pd.DataFrame()
|
|
self._data_str = pd.DataFrame()
|
|
self._horizontalHeaders = None
|
|
|
|
# Overloaded functions
|
|
|
|
def rowCount(self, parent=None):
|
|
return len(self._data_str.values)
|
|
|
|
def columnCount(self, parent=None):
|
|
return self._data_str.columns.size
|
|
|
|
def data(self, index, role=Qt.ItemDataRole.DisplayRole):
|
|
if not index.isValid():
|
|
return QtCore.QVariant()
|
|
|
|
if (role != Qt.ItemDataRole.DisplayRole) and (
|
|
role != Qt.ItemDataRole.EditRole):
|
|
return QtCore.QVariant()
|
|
|
|
if role == Qt.ItemDataRole.DisplayRole:
|
|
item = self._data_str.iloc[index.row(), index.column()]
|
|
return QtCore.QVariant(item)
|
|
elif role == Qt.ItemDataRole.EditRole:
|
|
item = self._data.iloc[index.row(), index.column()]
|
|
if type(item) is numpy.float64:
|
|
return QtCore.QVariant(float(item))
|
|
else:
|
|
return QtCore.QVariant(item)
|
|
|
|
def headerData(self, section, orientation,
|
|
role=Qt.ItemDataRole.DisplayRole):
|
|
|
|
if not ((orientation == Qt.Orientation.Horizontal)
|
|
and (role == Qt.ItemDataRole.DisplayRole)):
|
|
return super().headerData(section, orientation, role)
|
|
|
|
return self._horizontalHeaders[section]
|
|
|
|
# Other functions
|
|
|
|
def set_dataframe(self, df):
|
|
self._data = df
|
|
self._data_str = _get_str_dataframe(df)
|
|
self._horizontalHeaders = list(df.columns)
|
|
|
|
self.layoutAboutToBeChanged.emit()
|
|
self.layoutChanged.emit()
|
|
|
|
def get_dataframe(self) -> pd.DataFrame:
|
|
return self._data
|