apol: Implement CSV export of SEToolsTableView.

Closes #54

setoolsgui/tableview.py

@@ -16,16 +16,31 @@
 # License along with SETools.  If not, see
 # <http://www.gnu.org/licenses/>.
 #
+import csv
 
 from PyQt5.QtCore import Qt
-from PyQt5.QtGui import QKeySequence
-from PyQt5.QtWidgets import QApplication, QTableView
+from PyQt5.QtGui import QKeySequence, QCursor
+from PyQt5.QtWidgets import QAction, QApplication, QFileDialog, QMenu, QTableView
 
 
 class SEToolsTableView(QTableView):
 
     """QTableView class extended for SETools use."""
 
+    def __init__(self, parent):
+        super(SEToolsTableView, self).__init__(parent)
+
+        # set up right-click context menu
+        self.save_csv_action = QAction("Save table to CSV...", self)
+        self.menu = QMenu(self)
+        self.menu.addAction(self.save_csv_action)
+
+        # connect signals
+        self.save_csv_action.triggered.connect(self.choose_csv_save_location)
+
+    def contextMenuEvent(self, event):
+        self.menu.popup(QCursor.pos())
+
     def event(self, e):
         if e == QKeySequence.Copy or e == QKeySequence.Cut:
             datamodel = self.model()
@@ -54,3 +69,38 @@ class SEToolsTableView(QTableView):
 
         else:
             return super(SEToolsTableView, self).event(e)
+
+    def choose_csv_save_location(self):
+        filename = QFileDialog.getSaveFileName(self, "Save to CSV", "table.csv",
+                                               "Comma Separated Values Spreadsheet (*.csv);;"
+                                               "All Files (*)")[0]
+
+        if filename:
+            self.save_csv(filename)
+
+    def save_csv(self, filename):
+        """Save the current table data to the specified CSV file."""
+
+        datamodel = self.model()
+        row_count = datamodel.rowCount()
+        col_count = datamodel.columnCount()
+
+        with open(filename, 'w') as fd:
+            writer = csv.writer(fd, quoting=csv.QUOTE_MINIMAL)
+
+            # write headers
+            csv_row = []
+            for col in range(col_count):
+                csv_row.append(datamodel.headerData(col, Qt.Horizontal, Qt.DisplayRole))
+
+            writer.writerow(csv_row)
+
+            # write data
+            for row in range(row_count):
+                csv_row = []
+
+                for col in range(col_count):
+                    index = datamodel.index(row, col)
+                    csv_row.append(datamodel.data(index, Qt.DisplayRole))
+
+                writer.writerow(csv_row)