class ReferenceDataDlg(QDialog): def __init__(self, table, title, parent=None): super(ReferenceDataDlg, self).__init__(parent) self.create_widgets(table) self.layout_widgets() self.create_connections() self.setWindowTitle( "Asset Manager - Edit {0} Reference Data".format(title)) def create_widgets(self, table): self.model = QSqlTableModel(self) self.model.setTable(table) self.model.setSort(NAME, Qt.AscendingOrder) self.model.setHeaderData(ID, Qt.Horizontal, QVariant("ID")) self.model.setHeaderData(NAME, Qt.Horizontal, QVariant("Name")) self.model.setHeaderData(DESCRIPTION, Qt.Horizontal, QVariant("Description")) self.model.select() self.view = QTableView() self.view.setModel(self.model) self.view.setSelectionMode(QTableView.SingleSelection) self.view.setSelectionBehavior(QTableView.SelectRows) self.view.setColumnHidden(ID, True) self.view.resizeColumnsToContents() self.addButton = QPushButton("&Add") self.deleteButton = QPushButton("&Delete") self.okButton = QPushButton("&OK") def layout_widgets(self): buttonLayout = QHBoxLayout() buttonLayout.addWidget(self.addButton) buttonLayout.addWidget(self.deleteButton) buttonLayout.addStretch() buttonLayout.addWidget(self.okButton) layout = QVBoxLayout() layout.addWidget(self.view) layout.addLayout(buttonLayout) self.setLayout(layout) def create_connections(self): self.addButton.clicked.connect(self.addRecord) self.deleteButton.clicked.connect(self.deleteRecord) self.okButton.clicked.connect(self.accept) def addRecord(self): row = self.model.rowCount() self.model.insertRow(row) index = self.model.index(row, NAME) self.view.setCurrentIndex(index) self.view.edit(index) def deleteRecord(self): index = self.view.currentIndex() if not index.isValid(): return record = self.model.record(index.row()) id = record.value(ID).toInt()[0] table = self.model.tableName() query = QSqlQuery() if table == "deps": query.exec_(QString("SELECT COUNT(*) FROM employee " "WHERE deo_id = %1").arg(id)) elif table == "cities": query.exec_(QString("SELECT COUNT(*) FROM employee " "WHERE city_id = %1").arg(id)) count = 0 if query.next(): count = query.value(0).toInt()[0] if count: QMessageBox.information(self, QString("Delete %1").arg(table), (QString("Cannot delete %1<br>" "from the %2 table because it is used by " "%3 records") .arg(record.value(NAME).toString()) .arg(table).arg(count))) return self.model.removeRow(index.row()) self.model.submitAll()