diff --git a/banking_breakdown/ui.py b/banking_breakdown/ui.py
index 12d491a..ce1d486 100644
--- a/banking_breakdown/ui.py
+++ b/banking_breakdown/ui.py
@@ -11,6 +11,11 @@ from PyQt6.QtWidgets import QMainWindow, QPushButton, QHBoxLayout, QLabel, \
     QMessageBox, QFileDialog
 
 
+#
+# PandasModel
+#
+
+
 class PandasModel(QtCore.QAbstractTableModel):
     def __init__(self, df: pd.DataFrame, parent=None):
         QtCore.QAbstractTableModel.__init__(self, parent)
@@ -66,6 +71,11 @@ class PandasModel(QtCore.QAbstractTableModel):
         return self._data
 
 
+#
+# MainWindow
+#
+
+
 class MainWindow(QMainWindow):
     def __init__(self):
         super(MainWindow, self).__init__()
@@ -98,25 +108,6 @@ class MainWindow(QMainWindow):
         header.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
         header.customContextMenuRequested.connect(self._header_right_clicked)
 
-    def _add_warning_text(self, text: str):
-        layout = QHBoxLayout()
-
-        warningIcon = QLabel()
-        pixmap = QPixmap("res/warning.png")
-        warningIcon.setPixmap(pixmap)
-
-        label = QLabel(text)
-        label.setWordWrap(True)
-
-        layout.addWidget(warningIcon)
-        layout.addWidget(label)
-
-        layout.setStretch(0, 0)
-        layout.setStretch(1, 1)
-
-        self._warning_layout.addLayout(layout)
-        self._warnings.append((layout, warningIcon, label))
-
     def set_statement_data(self, df: pd.DataFrame):
         if 'category' not in df.columns:
             df['category'] = [' '] * len(df.index)
@@ -138,6 +129,32 @@ class MainWindow(QMainWindow):
 
         self._show_warnings(df)
 
+    def get_statement_data(self) -> pd.DataFrame:
+        return self._table_view.model().sourceModel().get_dataframe()
+
+    def add_categories(self, categories: typing.Sequence[str]):
+        for category in categories:
+            self._list_widget.addItem(category)
+
+    def _add_warning_text(self, text: str):
+        layout = QHBoxLayout()
+
+        warningIcon = QLabel()
+        pixmap = QPixmap("res/warning.png")
+        warningIcon.setPixmap(pixmap)
+
+        label = QLabel(text)
+        label.setWordWrap(True)
+
+        layout.addWidget(warningIcon)
+        layout.addWidget(label)
+
+        layout.setStretch(0, 0)
+        layout.setStretch(1, 1)
+
+        self._warning_layout.addLayout(layout)
+        self._warnings.append((layout, warningIcon, label))
+
     def _show_warnings(self, df: pd.DataFrame):
         for (layout, icon, text) in self._warnings:
             icon.hide()
@@ -162,13 +179,6 @@ class MainWindow(QMainWindow):
                 " rename the column containing the balance"
                 " after each transaction to 'balance'")
 
-    def get_statement_data(self) -> pd.DataFrame:
-        return self._table_view.model().sourceModel().get_dataframe()
-
-    def add_categories(self, categories: typing.Sequence[str]):
-        for category in categories:
-            self._list_widget.addItem(category)
-
     def _header_right_clicked(self, pos):
         context = QMenu(self)