Proper element creation and deletion

1 files changed, 96 insertions, 17 deletions

diff --git a/src/ui_table.py b/src/ui_table.py
index 7640ace..db69beb 100644
--- a/src/ui_table.py
+++ b/src/ui_table.py
@@ -3,6 +3,7 @@ from PyQt5.QtWidgets import *
 from PyQt5.QtGui import *
 from PyQt5.QtCore import *
 
+import phys
 from ui import GUI, Pens
 from ui_widgets import *
 
@@ -12,8 +13,8 @@ class ElementEditorWindow(QMainWindow):
 	windows = list()
 
 	@staticmethod
-	def open(system, pol):
-		win = ElementEditorWindow(system, pol)
+	def open(row):
+		win = ElementEditorWindow(row)
 		win.show()
 		ElementEditorWindow.windows.append(win)
 
@@ -24,10 +25,11 @@ class ElementEditorWindow(QMainWindow):
 				win.close()
 				del(win)
 
-	def __init__(self, system, pol):
+	def __init__(self, row):
 		super().__init__(flags=Qt.Dialog)
-		self.system = system
-		self.pol = pol
+		self.row = row
+		self.system = row.parent.system
+		self.pol = row.pol
 
 		self.setWindowTitle("Element editor")
 		root = QVBoxLayout()
@@ -50,7 +52,7 @@ class ElementEditorWindow(QMainWindow):
 		self.type = QComboBox()
 		self.type.addItem("Linear", "linear",)
 		self.type.addItem("Quarterwave plate", "quarterwave")
-		self.type.setCurrentIndex(0 if pol.type == "linear" else 1)
+		self.type.setCurrentIndex(0 if self.pol.type == "linear" else 1)
 		self.type.currentIndexChanged.connect(self.change_type)
 		root.addWidget(self.type)
 
@@ -91,7 +93,7 @@ class ElementEditorWindow(QMainWindow):
 
 	def change_name(self):
 		self.pol.name = self.name.text()
-		self.pol.row.optbox.name.setText(self.pol.name)
+		self.row.optbox.name.setText(self.pol.name)
 
 	def change_type(self):
 		self.pol.set_type(self.type.currentData())
@@ -128,8 +130,6 @@ class OptBox(QVBoxLayout):
 
 		# Name
 		self.name = MeinGroßLabel(pol.name)
-		self.name.setContextMenuPolicy(Qt.CustomContextMenu)
-		self.name.customContextMenuRequested.connect(lambda pos: OptBox.open_menu(self, pos))
 		self.addWidget(self.name)
 
 		# Enable checkbox
@@ -145,14 +145,6 @@ class OptBox(QVBoxLayout):
 		self.addLayout(self.angle)
 		self.addItem(ExpandingSpacer())
 
-	def open_menu(self, pos):
-		menu = QMenu()
-		menu.addAction("Edit", lambda: ElementEditorWindow.open(self.row.parent.system, self.row.pol))
-		menu.addAction("Insert before", lambda: half_assed_element_creation(self.rownum - 1))
-		menu.addAction("Insert after", lambda: half_assed_element_creation(self.rownum))
-		menu.addAction("Delete", lambda: half_assed_element_deletion(self.rownum - 1))
-		menu.exec_(self.name.mapToGlobal(pos))
-
 	def change_enable(self):
 		self.row.pol.enable = self.enable.isChecked()
 		update()
@@ -211,10 +203,97 @@ class SystemTable(QTableWidget):
 		self.vh = self.verticalHeader()
 		self.vh.setSectionResizeMode(QHeaderView.ResizeToContents)
 
+		# Smooth scrolling
+		self.setHorizontalScrollMode(QAbstractItemView.ScrollPerPixel)
+		self.setVerticalScrollMode(QAbstractItemView.ScrollPerPixel)
+
+		# Select only entire rows
+		self.setSelectionBehavior(QAbstractItemView.SelectRows)
+	
+		self.setContextMenuPolicy(Qt.CustomContextMenu)
+		self.customContextMenuRequested.connect(self.context_menu)
+
+	def context_menu(self, pos):
+		indices = list()
+		for lo, hi in [(r.topRow(), r.bottomRow()) \
+		               for r in self.selectedRanges()]:
+			indices.extend(list(range(lo, hi + 1)))
+
+		menu = QMenu()
+
+		if len(indices) == 0:
+			if len(self.rows) > 0:
+				menu.addAction("Add a new element at the top", \
+				               lambda: self.insert_row(None, True))
+				menu.addAction("Add a new element at the bottom", \
+				               lambda: self.insert_row(None, False))
+			else:
+				menu.addAction("Add a new element", lambda: self.insert_row(None, False))
+		elif len(indices) == 1:
+			i = indices[0]
+			menu.addAction("Edit %s" % self.rows[i].pol.name, \
+			               lambda: ElementEditorWindow.open(self.rows[i]))
+			menu.addSeparator()
+			menu.addAction("Add a new element before", lambda: self.insert_row(self.rows[i], True))
+			menu.addAction("Add a new element after", lambda: self.insert_row(self.rows[i], False))
+			deltext = "Delete %s" % self.rows[i].pol.name
+		else:
+			deltext = "Delete %d elements" % len(indices)
+
+		if len(indices) >= 1:
+			menu.addAction(deltext, lambda: GUI.table.delete_rows(indices))
+
+		menu.addSeparator()
+
+		menu.exec_(self.mapToGlobal(pos))
+
+	def clear(self):
+		ElementEditorWindow.close_all()
+		system.elements = list()
+		self.rows = list()
+		self.clearContents()
+		GUI.do_update()
+
 	def renumber_rows(self):
 		for i, row in enumerate(self.rows):
 			row.row_number = i
 
+	def insert_row(self, wrt, before, pol=None):
+		if wrt is not None:
+			i = wrt.row_number + 0 if before else 1
+		else:
+			i = 0 if before else len(self.rows)
+
+		if pol is None:
+			pol = phys.Polarizer("linear")
+			editor = True
+		else:
+			editor = False
+
+		self.system.elements.insert(i, pol)
+		self.insertRow(i)
+		row = SystemTableRow(self, pol, i)
+		self.rows.insert(i, row)
+		self.renumber_rows()
+		GUI.do_update()
+
+		if editor:
+			ElementEditorWindow.open(row)
+
+	def delete_rows(self, indices):
+		for i in sorted(indices):
+			ElementEditorWindow.close_all(pol=self.rows[i].pol)
+
+		offs = 0
+		for i in sorted(indices):
+			self.removeRow(i - offs)
+			del(self.rows[i - offs])
+			del(self.system.elements[i - offs])
+			offs += 1
+	
+		self.renumber_rows()
+		GUI.do_update()
+
 	def populate(self, system):
 		self.system = system