Proper element creation and deletion

2 files changed, 133 insertions, 90 deletions

diff --git a/src/ui.py b/src/ui.py
index 0077b13..fc85dbd 100644
--- a/src/ui.py
+++ b/src/ui.py
@@ -14,11 +14,11 @@ class GUI:
 	monospace = QFont("monospace")
 	bigfont = QFont("sans-serif", pointSize=10, weight=1000)
 
-
-
 from ui_widgets import *
 from ui_table import *
 
+
+
 def update():
 	for i, pol in enumerate(system.elements):
 		row = GUI.table.rows[i]
@@ -73,40 +73,6 @@ def update():
 
 GUI.do_update = update
 
-
-def half_assed_element_creation(index=None):
-	pol = phys.Polarizer("linear")
-	if index is None:
-		system.elements.append(pol)
-	else:
-		system.elements.insert(index, pol)
-
-	populate_table()
-	GUI.table_frame = QFrame()
-	GUI.table_frame.setLayout(GUI.table)
-	GUI.scroll.setWidget(GUI.table_frame)
-
-	ElementEditorWindow.open(pol)
-	update()
-
-def half_assed_clear():
-	ElementEditorWindow.close_all()
-	system.elements = list()
-	populate_table()
-	GUI.table_frame = QFrame()
-	GUI.table_frame.setLayout(GUI.table)
-	GUI.scroll.setWidget(GUI.table_frame)
-	update()
-
-def half_assed_element_deletion(index):
-	ElementEditorWindow.close_all(pol=system.elements[pol])
-	populate_table()
-	GUI.table_frame = QFrame()
-	GUI.table_frame.setLayout(GUI.table)
-	GUI.scroll.setWidget(GUI.table_frame)
-	update()
-
-
 # FIXME: refactor
 #def optimize(which):
 #	if len(system.elements) == 0:
@@ -144,12 +110,42 @@ class MainWindow(QMainWindow):
 	def __init__(self, system_):
 		super().__init__()
 		self.last_save_path = None
+		global system
+		system = system_
 
-		self.statusBar()
 		self.setWindowTitle("Polarizzazione italiana")
 
-		global system
-		system = system_
+		# Central widget
+		root = QVBoxLayout()
+		self.setCentralWidget(LayoutWrapper(root))
+
+		box = QHBoxLayout()
+		root.addLayout(box)
+
+		# Top bar
+		box.addWidget(QLabel("Input intensity"))
+		GUI.input_intensity = QLineEdit("%g" % system.input_intensity)
+		GUI.input_intensity.textChanged.connect(self.change_input_intensity)
+		box.addWidget(GUI.input_intensity)
+
+		# Splitter
+		split = QSplitter()
+		root.addWidget(split)
+
+		# Table (LHS)
+		GUI.table = SystemTable()
+		split.addWidget(GUI.table)
+
+		# Splitter RHS
+		rhs = QVBoxLayout()
+		split.addWidget(LayoutWrapper(rhs))
+
+		Widocques.image = QImage("jones.jpg")
+		GUI.widok = Widocques()
+		rhs.addWidget(GUI.widok)
+
+		GUI.table.populate(system)
+		update()
 
 		# Menu
 		menu_bar = self.menuBar()
@@ -182,45 +178,13 @@ class MainWindow(QMainWindow):
 
 		act = QAction("&Add a new element", self)
 		act.setShortcut("Ctrl+N")
-		act.triggered.connect(lambda: half_assed_element_creation())
+		act.triggered.connect(lambda: GUI.table.insert_row(None, False))
 		menu.addAction(act)
 
 		act = QAction("&Remove all elements", self)
-		act.triggered.connect(half_assed_clear)
+		act.triggered.connect(GUI.table.clear)
 		menu.addAction(act)
 
-		# Central widget
-		root = QVBoxLayout()
-		self.setCentralWidget(LayoutWrapper(root))
-
-		box = QHBoxLayout()
-		root.addLayout(box)
-
-		# Top bar
-		box.addWidget(QLabel("Input intensity"))
-		GUI.input_intensity = QLineEdit("%g" % system.input_intensity)
-		GUI.input_intensity.textChanged.connect(self.change_input_intensity)
-		box.addWidget(GUI.input_intensity)
-
-		# Splitter
-		split = QSplitter()
-		root.addWidget(split)
-
-		# Table (LHS)
-		GUI.table = SystemTable()
-		split.addWidget(GUI.table)
-
-		# Splitter RHS
-		rhs = QVBoxLayout()
-		split.addWidget(LayoutWrapper(rhs))
-
-		Widocques.image = QImage("jones.jpg")
-		GUI.widok = Widocques()
-		rhs.addWidget(GUI.widok)
-
-		GUI.table.populate(system)
-		update()
-
 	def change_input_intensity(self):
 		try:
 			system.input_intensity = float(GUI.input_intensity.text())
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