159 lines
5.0 KiB
Python
159 lines
5.0 KiB
Python
from PySide6.QtCore import QAbstractTableModel, Qt
|
|
from PySide6.QtGui import QBrush, QColor
|
|
|
|
from aare.common.models import SampleShortInfo, SampleShortInfoList
|
|
|
|
|
|
def get_entry(sample: SampleShortInfo, column: int, *, show_user: bool = False):
|
|
if show_user:
|
|
if column == 0:
|
|
return sample.user
|
|
elif column == 1:
|
|
return sample.dewar_name
|
|
elif column == 2:
|
|
return sample.puck_name
|
|
elif column == 3:
|
|
return sample.sample_name
|
|
else:
|
|
if column == 0:
|
|
return sample.dewar_name
|
|
elif column == 1:
|
|
return sample.puck_name
|
|
elif column == 2:
|
|
return sample.sample_name
|
|
return ""
|
|
|
|
|
|
class SampleQueueSpreadsheet(QAbstractTableModel):
|
|
def __init__(
|
|
self,
|
|
parent=None,
|
|
samples: list[SampleShortInfo] | None = None,
|
|
show_user: bool = False,
|
|
):
|
|
super().__init__(parent)
|
|
if samples is None:
|
|
samples = []
|
|
self.__running = False
|
|
self._show_user = bool(show_user)
|
|
self.samples: list[SampleShortInfo] = samples
|
|
self.header = (
|
|
["User", "Dewar", "Puck", "Sample"]
|
|
if self._show_user
|
|
else ["Dewar", "Puck", "Sample"]
|
|
)
|
|
|
|
def flags(self, index):
|
|
default_flags = super().flags(index)
|
|
return Qt.ItemFlag.ItemIsDropEnabled | default_flags
|
|
|
|
def rowCount(self, parent=None):
|
|
return len(self.samples)
|
|
|
|
def columnCount(self, parent=None):
|
|
return len(self.header)
|
|
|
|
def data(self, index, role=None):
|
|
if role == Qt.ItemDataRole.DisplayRole:
|
|
return get_entry(self.samples[index.row()], index.column(), show_user=self._show_user)
|
|
elif role == Qt.ItemDataRole.TextAlignmentRole:
|
|
return Qt.AlignmentFlag.AlignCenter
|
|
elif role == Qt.ItemDataRole.BackgroundRole:
|
|
if index.row() == 0:
|
|
if self.__running:
|
|
return QBrush(QColor(255, 102, 0))
|
|
else:
|
|
return QBrush(QColor(114, 159, 207))
|
|
return QBrush(QColor(255, 255, 255))
|
|
return None
|
|
|
|
def headerData(self, section, orientation, role=None):
|
|
if role == Qt.ItemDataRole.DisplayRole:
|
|
if orientation == Qt.Orientation.Horizontal:
|
|
return self.header[section] if self.header else f"Column {section + 1}"
|
|
if orientation == Qt.Orientation.Vertical:
|
|
return str(section + 1)
|
|
return None
|
|
|
|
def updateData(self, samples: list[SampleShortInfo],):
|
|
self.beginResetModel()
|
|
self.samples = list(samples)
|
|
self.endResetModel()
|
|
|
|
def mimeTypes(self):
|
|
return ['text/plain']
|
|
|
|
def canDropMimeData(self, data, action, row, column, parent):
|
|
if data.hasText():
|
|
return True
|
|
return False
|
|
|
|
def dropMimeData(self, data, action, row, column, parent):
|
|
if not self.canDropMimeData(data, action, row, column, parent):
|
|
return False
|
|
|
|
if parent.isValid():
|
|
row = parent.row()
|
|
|
|
try:
|
|
l = SampleShortInfoList.model_validate_json(data.text())
|
|
self.beginResetModel()
|
|
for sample in l.s:
|
|
updated_row = row
|
|
updated_samples = []
|
|
for i in range(len(self.samples)):
|
|
if sample != self.samples[i]:
|
|
updated_samples.append(self.samples[i])
|
|
else:
|
|
if 0 < row < i:
|
|
updated_row -= 1
|
|
if updated_row == -1:
|
|
updated_samples.append(sample)
|
|
else:
|
|
updated_samples.insert(updated_row, sample)
|
|
self.samples = updated_samples
|
|
self.endResetModel()
|
|
except Exception as e:
|
|
print(e)
|
|
return False
|
|
return True
|
|
|
|
def set_running(self, running: bool):
|
|
self.__running = running
|
|
self.layoutChanged.emit()
|
|
|
|
def remove_sample(self, db_id: int):
|
|
self.beginResetModel()
|
|
self.samples = list(filter(lambda x: x.db_id != db_id, self.samples))
|
|
self.endResetModel()
|
|
|
|
def clearSamples(self):
|
|
self.beginResetModel()
|
|
self.samples = []
|
|
self.endResetModel()
|
|
|
|
def annotate_sample(self, db_id: int, comment: str) -> None:
|
|
self.beginResetModel()
|
|
updated_samples: list[SampleShortInfo] = []
|
|
for sample in self.samples:
|
|
if sample.db_id == db_id:
|
|
updated_samples.append(sample.model_copy(update={"comment": comment}))
|
|
else:
|
|
updated_samples.append(sample)
|
|
self.samples = updated_samples
|
|
self.endResetModel()
|
|
|
|
def to_state(self) -> dict:
|
|
return {
|
|
"samples": [s.model_dump() for s in self.samples],
|
|
}
|
|
|
|
def from_state(self, state: dict):
|
|
self.beginResetModel()
|
|
|
|
self.samples = [
|
|
SampleShortInfo.model_validate(d)
|
|
for d in state.get("samples", [])
|
|
]
|
|
|
|
self.endResetModel() |