diff --git a/b_asic/GUI/__init__.py b/b_asic/GUI/__init__.py
index 61d8467332aee05a21d2ac7d4e9c96db04020bf2..357ada54290822d2b6c92c4e831446e64de6a908 100644
--- a/b_asic/GUI/__init__.py
+++ b/b_asic/GUI/__init__.py
@@ -2,6 +2,6 @@
 
 Graphical user interface for B-ASIC.
 """
-from b_asic.GUI.main_window import MainWindow, start_editor
+from b_asic.GUI.main_window import start_editor
 
-__all__ = ['MainWindow', 'start_editor']
+__all__ = ['start_editor']
diff --git a/b_asic/GUI/arrow.py b/b_asic/GUI/arrow.py
index 2bca3cf74d908c153ac05df615b715213e832d8d..e62a9cbe936d2f8b4ad5255e87149136693d17a1 100644
--- a/b_asic/GUI/arrow.py
+++ b/b_asic/GUI/arrow.py
@@ -1,3 +1,5 @@
+from typing import TYPE_CHECKING, Optional, cast
+
 from qtpy.QtCore import QPointF
 from qtpy.QtGui import QPainterPath, QPen
 from qtpy.QtWidgets import QGraphicsPathItem, QMenu
@@ -5,35 +7,49 @@ from qtpy.QtWidgets import QGraphicsPathItem, QMenu
 from b_asic.GUI._preferences import GRID, LINECOLOR, PORTHEIGHT, PORTWIDTH
 from b_asic.signal import Signal
 
+if TYPE_CHECKING:
+    from b_asic.GUI.drag_button import DragButton
+    from b_asic.GUI.main_window import SFGMainWindow
+    from b_asic.GUI.port_button import PortButton
+    from b_asic.port import InputPort, OutputPort
+
 
 class Arrow(QGraphicsPathItem):
-    """Arrow/connection in signal flow graph GUI."""
+    """
+    Arrow/connection in signal flow graph GUI.
 
-    def __init__(self, source, destination, window, signal=None, parent=None):
-        """
-        Parameters
-        ----------
-        source
-            Source operation.
-        destination
-            Destination operation.
-        window
-            Window containing signal flow graph.
-        signal : Signal, optional
-            Let arrow represent *signal*.
-        parent : optional
-            Parent.
-        """
+    Parameters
+    ----------
+    source_port_button : :class:`~b_asic.GUI.port_button.PortButton`
+        Source port button.
+    destination_port_button : :class:`~b_asic.GUI.port_button.PortButton`
+        Destination port button.
+    window : :class:`~b_asic.GUI.main_window.SFGMainWindow`
+        Window containing signal flow graph.
+    signal : Signal, optional
+        Let arrow represent *signal*.
+    parent : optional
+        Parent.
+    """
+
+    def __init__(
+        self,
+        source_port_button: "PortButton",
+        destination_port_button: "PortButton",
+        window: "SFGMainWindow",
+        signal: Optional[Signal] = None,
+        parent=None,
+    ):
         super().__init__(parent)
-        self.source = source
+        self._source_port_button = source_port_button
         if signal is None:
-            signal = Signal(source.port, destination.port)
+            signal = Signal(source_port_button.port, destination_port_button.port)
         self.signal = signal
-        self.destination = destination
+        self._destination_port_button = destination_port_button
         self._window = window
-        self.moveLine()
-        self.source.moved.connect(self.moveLine)
-        self.destination.moved.connect(self.moveLine)
+        self.update_arrow()
+        self._source_port_button.moved.connect(self.update_arrow)
+        self._destination_port_button.moved.connect(self.update_arrow)
 
     def contextMenuEvent(self, event):
         """Open right-click menu."""
@@ -41,16 +57,62 @@ class Arrow(QGraphicsPathItem):
         menu.addAction("Delete", self.remove)
         menu.exec_(self.cursor().pos())
 
+    @property
+    def source_operation_button(self) -> "DragButton":
+        """The source DragButton."""
+        return self._source_port_button._operation_button
+
+    @property
+    def destination_operation_button(self) -> "DragButton":
+        """The destination DragButton."""
+        return self._destination_port_button._operation_button
+
+    @property
+    def source_operation(self) -> "Operation":
+        """The source Operation."""
+        return self._source_port_button.operation
+
+    @property
+    def destination_operation(self) -> "Operation":
+        """The destination Operation."""
+        return self._destination_port_button.operation
+
+    @property
+    def source_port_button(self) -> "PortButton":
+        """The source PortButton."""
+        return self._source_port_button
+
+    @property
+    def destination_port_button(self) -> "PortButton":
+        """The destination PortButton."""
+        return self._destination_port_button
+
+    @property
+    def source_port(self) -> "OutputPort":
+        """The source OutputPort."""
+        return cast("OutputPort", self._source_port_button.port)
+
+    @property
+    def desination_port(self) -> "InputPort":
+        """The destination InputPort."""
+        return cast("InputPort", self._destination_port_button.port)
+
+    def set_source_operation(self, source: "Operation"):
+        """Set operation of the source DragButton"""
+        self._source_port_button._operation_button.operation = source
+
+    def set_destination_operation(self, destination: "Operation"):
+        """Set operation of the destination DragButton"""
+        self._destination_port_button._operation_button.operation = destination
+
     def remove(self):
         """Remove line and connections to signals etc."""
         self.signal.remove_destination()
         self.signal.remove_source()
         self._window._scene.removeItem(self)
-        if self in self._window._arrows:
-            self._window._arrows.remove(self)
 
-        if self in self._window._signal_ports:
-            for port1, port2 in self._window._signal_ports[self]:
+        if self in self._window._arrow_ports:
+            for port1, port2 in self._window._arrow_ports[self]:
                 for (
                     operation,
                     operation_ports,
@@ -73,30 +135,39 @@ class Arrow(QGraphicsPathItem):
                             is not self._window._operation_to_sfg[operation]
                         }
 
-        del self._window._signal_ports[self]
+        del self._window._arrow_ports[self]
 
-    def moveLine(self):
+    def update_arrow(self):
         """
-        Draw a line connecting ``self.source`` with ``self.destination``.
+        Update coordinates for arrow.
+
         Used as callback when moving operations.
         """
         ORTHOGONAL = True
         OFFSET = 2 * PORTWIDTH
         self.setPen(QPen(LINECOLOR, 3))
-        source_flipped = self.source.operation.is_flipped()
-        destination_flipped = self.destination.operation.is_flipped()
+        source_flipped = self.source_operation_button.is_flipped()
+        destination_flipped = self.destination_operation_button.is_flipped()
         x0 = (
-            self.source.operation.x()
-            + self.source.x()
+            self.source_operation_button.x()
+            + self.source_port_button.x()
             + (PORTWIDTH if not source_flipped else 0)
         )
-        y0 = self.source.operation.y() + self.source.y() + PORTHEIGHT / 2
+        y0 = (
+            self.source_operation_button.y()
+            + self.source_port_button.y()
+            + PORTHEIGHT / 2
+        )
         x1 = (
-            self.destination.operation.x()
-            + self.destination.x()
+            self.destination_operation_button.x()
+            + self.destination_port_button.x()
             + (0 if not destination_flipped else PORTWIDTH)
         )
-        y1 = self.destination.operation.y() + self.destination.y() + PORTHEIGHT / 2
+        y1 = (
+            self.destination_operation_button.y()
+            + self._destination_port_button.y()
+            + PORTHEIGHT / 2
+        )
         xmid = (x0 + x1) / 2
         ymid = (y0 + y1) / 2
         both_flipped = source_flipped and destination_flipped
diff --git a/b_asic/GUI/drag_button.py b/b_asic/GUI/drag_button.py
index eeb63255be109379f0273a57664745ed66f91046..7ed8c36119945460e9253e24ba493c570e6aa8de 100644
--- a/b_asic/GUI/drag_button.py
+++ b/b_asic/GUI/drag_button.py
@@ -5,7 +5,7 @@ Contains a GUI class for drag buttons.
 """
 
 import os.path
-from typing import List
+from typing import TYPE_CHECKING, List
 
 from qtpy.QtCore import QSize, Qt, Signal
 from qtpy.QtGui import QIcon
@@ -18,6 +18,9 @@ from b_asic.gui_utils.decorators import decorate_class, handle_error
 from b_asic.operation import Operation
 from b_asic.port import InputPort
 
+if TYPE_CHECKING:
+    from b_asic.GUI.main_window import SFGMainWindow
+
 
 @decorate_class(handle_error)
 class DragButton(QPushButton):
@@ -29,8 +32,11 @@ class DragButton(QPushButton):
     Parameters
     ----------
     operation : :class:`~b_asic.operation.Operation`
-    is_show_name : bool
-    window
+        The operation that the drag button corresponds to.
+    show_name : bool
+        Whether to show the name.
+    window : SFGMainWindow
+        Parent MainWindow.
     parent
     """
 
@@ -40,16 +46,15 @@ class DragButton(QPushButton):
     def __init__(
         self,
         operation: Operation,
-        is_show_name: bool,
+        show_name: bool,
         window,
         parent=None,
     ):
         self.name = operation.graph_id
-        self.ports: List[PortButton] = []
-        self.is_show_name = is_show_name
+        self._ports: List[PortButton] = []
+        self.show_name = show_name
         self._window = window
         self.operation = operation
-        self.clicked = 0
         self.pressed = False
         self._m_press = False
         self._m_drag = False
@@ -78,11 +83,20 @@ class DragButton(QPushButton):
             self._context_menu = menu
         self._context_menu.exec_(self.cursor().pos())
 
-    def show_properties_window(self, event=None):
+    def show_properties_window(self, event=None) -> None:
+        """Display the properties window for the associated Operation."""
         self._properties_window = PropertiesWindow(self, self._window)
         self._properties_window.show()
 
-    def add_label(self, label):
+    def add_label(self, label: str) -> None:
+        """
+        Add label to button.
+
+        Parameters
+        ----------
+        label : src
+            The label to add.
+        """
         self.label = label
 
     def mousePressEvent(self, event):
@@ -94,6 +108,11 @@ class DragButton(QPushButton):
 
         super().mousePressEvent(event)
 
+    @property
+    def port_list(self) -> List[PortButton]:
+        """Return a list of PortButtons."""
+        return self._ports
+
     def mouseMoveEvent(self, event):
         if event.buttons() == Qt.MouseButton.LeftButton and self._m_press:
             self._m_drag = True
@@ -133,7 +152,7 @@ class DragButton(QPushButton):
 
     def _flip(self, event=None):
         self._flipped = not self._flipped
-        for pb in self.ports:
+        for pb in self._ports:
             if isinstance(pb.port, InputPort):
                 newx = MINBUTTONSIZE - PORTWIDTH if self._flipped else 0
             else:
@@ -164,7 +183,15 @@ class DragButton(QPushButton):
         """Return True if the button is flipped (inputs to the right)."""
         return self._flipped
 
-    def select_button(self, modifiers=None):
+    def select_button(self, modifiers=None) -> None:
+        """
+        Select button taking *modifiers* into account.
+
+        Parameters
+        ----------
+        modifiers : optional
+            Qt keyboard modifier.
+        """
         if modifiers != Qt.KeyboardModifier.ControlModifier:
             for button in self._window._pressed_operations:
                 button._toggle_button(button.pressed)
@@ -179,8 +206,8 @@ class DragButton(QPushButton):
             else:
                 self._window._pressed_operations.append(self)
 
-        for signal in self._window._arrows:
-            signal.update()
+        for arrow in self._window._arrow_ports:
+            arrow.update()
 
     def remove(self, event=None):
         """Remove button/operation from signal flow graph."""
@@ -188,10 +215,10 @@ class DragButton(QPushButton):
         self._window._scene.removeItem(self._window._drag_operation_scenes[self])
 
         _signals = []
-        for signal, ports in self._window._signal_ports.items():
+        for signal, ports in self._window._arrow_ports.items():
             if any(
                 map(
-                    lambda port: set(port).intersection(set(self.ports)),
+                    lambda port: set(port).intersection(set(self._ports)),
                     ports,
                 )
             ):
@@ -217,7 +244,7 @@ class DragButton(QPushButton):
                 is not self._window._operation_to_sfg[self]
             }
 
-        for port in self._window._ports[self]:
+        for port in self._ports:
             if port in self._window._pressed_ports:
                 self._window._pressed_ports.remove(port)
 
@@ -231,6 +258,8 @@ class DragButton(QPushButton):
             del self._window._drag_buttons[self.operation]
 
     def add_ports(self):
+        """Add ports to button."""
+
         def _determine_port_distance(opheight, ports):
             """
             Determine the distance between each port on the side of an operation.
@@ -251,11 +280,11 @@ class DragButton(QPushButton):
             port.setFixedSize(PORTWIDTH, PORTHEIGHT)
             port.move(0, dist)
             port.show()
-            self.ports.append(port)
+            self._ports.append(port)
 
         for i, dist in enumerate(output_ports_dist):
             port = PortButton(">", self, op.output(i))
             port.setFixedSize(PORTWIDTH, PORTHEIGHT)
             port.move(MINBUTTONSIZE - PORTWIDTH, dist)
             port.show()
-            self.ports.append(port)
+            self._ports.append(port)
diff --git a/b_asic/GUI/main_window.py b/b_asic/GUI/main_window.py
index bdb3f805f9209228c44ecfb16438acf20e46b373..414d0084a72d3bb97ba2630c4b0a96cc0a6d89c2 100644
--- a/b_asic/GUI/main_window.py
+++ b/b_asic/GUI/main_window.py
@@ -1,15 +1,16 @@
-"""B-ASIC Main Window Module.
+"""
+B-ASIC Signal Flow Graph Editor Module.
 
-This file opens the main window of the GUI for B-ASIC when run.
+This file opens the main SFG editor window of the GUI for B-ASIC when run.
 """
 
 
-import importlib
+import importlib.util
 import logging
 import os
 import sys
 from collections import deque
-from typing import Dict, List, Optional, Tuple, cast
+from typing import TYPE_CHECKING, Deque, Dict, List, Optional, Sequence, Tuple, cast
 
 from qtpy.QtCore import QCoreApplication, QFileInfo, QSettings, QSize, Qt, QThread
 from qtpy.QtGui import QCursor, QIcon, QKeySequence, QPainter
@@ -52,6 +53,9 @@ from b_asic.signal_flow_graph import SFG
 from b_asic.simulation import Simulation
 from b_asic.special_operations import Input, Output
 
+if TYPE_CHECKING:
+    from qtpy.QtWidgets import QGraphicsProxyWidget
+
 logging.basicConfig(level=logging.INFO)
 
 QCoreApplication.setOrganizationName("Linköping University")
@@ -61,31 +65,29 @@ QCoreApplication.setApplicationVersion(__version__)
 
 
 @decorate_class(handle_error)
-class MainWindow(QMainWindow):
+class SFGMainWindow(QMainWindow):
     def __init__(self):
         super().__init__()
         self._ui = Ui_main_window()
         self._ui.setupUi(self)
         self.setWindowIcon(QIcon("small_logo.png"))
         self._scene = QGraphicsScene(self)
-        self._operations_from_name = {}
+        self._operations_from_name: Dict[str, Operation] = {}
         self._zoom = 1
-        self._sfg_name_i = 0
-        self._drag_operation_scenes = {}
+        self._drag_operation_scenes: Dict[DragButton, "QGraphicsProxyWidget"] = {}
         self._drag_buttons: Dict[Operation, DragButton] = {}
-        self._arrows: List[Arrow] = []
         self._mouse_pressed = False
         self._mouse_dragging = False
         self._starting_port = None
-        self._pressed_operations = []
-        self._ports = {}
-        self._signal_ports = {}
+        self._pressed_operations: List[DragButton] = []
+        self._arrow_ports: Dict[Arrow, List[Tuple[PortButton, PortButton]]] = {}
         self._operation_to_sfg: Dict[DragButton, SFG] = {}
-        self._pressed_ports = []
-        self._sfg_dict = {}
+        self._pressed_ports: List[PortButton] = []
+        self._sfg_dict: Dict[str, SFG] = {}
         self._window = self
         self._logger = logging.getLogger(__name__)
-        self._plot: Dict[Simulation, PlotWindow] = dict()
+        self._plot: Dict[Simulation, PlotWindow] = {}
+        self._ports: Dict[DragButton, List[PortButton]] = {}
 
         # Create Graphics View
         self._graphics_view = QGraphicsView(self._scene, self)
@@ -102,8 +104,8 @@ class MainWindow(QMainWindow):
 
         # Add operations
         self._max_recent_files = 4
-        self._recent_files = []
-        self._recent_files_paths = deque(maxlen=self._max_recent_files)
+        self._recent_files_actions: List[QAction] = []
+        self._recent_files_paths: Deque[str] = deque(maxlen=self._max_recent_files)
 
         self.add_operations_from_namespace(
             b_asic.core_operations, self._ui.core_operations_list
@@ -112,14 +114,16 @@ class MainWindow(QMainWindow):
             b_asic.special_operations, self._ui.special_operations_list
         )
 
-        self._shortcut_core = QShortcut(QKeySequence("Ctrl+R"), self._ui.operation_box)
-        self._shortcut_core.activated.connect(
+        self._shortcut_refresh_operations = QShortcut(
+            QKeySequence("Ctrl+R"), self._ui.operation_box
+        )
+        self._shortcut_refresh_operations.activated.connect(
             self._refresh_operations_list_from_namespace
         )
         self._scene.selectionChanged.connect(self._select_operations)
 
         self.move_button_index = 0
-        self.is_show_names = True
+        self._show_names = True
 
         self._check_show_names = QAction("Show operation names")
         self._check_show_names.triggered.connect(self.view_operation_names)
@@ -133,13 +137,13 @@ class MainWindow(QMainWindow):
         self._ui.aboutBASIC.triggered.connect(self.display_about_page)
         self._ui.keybindsBASIC.triggered.connect(self.display_keybindings_page)
         self._ui.core_operations_list.itemClicked.connect(
-            self.on_list_widget_item_clicked
+            self._on_list_widget_item_clicked
         )
         self._ui.special_operations_list.itemClicked.connect(
-            self.on_list_widget_item_clicked
+            self._on_list_widget_item_clicked
         )
         self._ui.custom_operations_list.itemClicked.connect(
-            self.on_list_widget_item_clicked
+            self._on_list_widget_item_clicked
         )
         self._ui.save_menu.triggered.connect(self.save_work)
         self._ui.load_menu.triggered.connect(self.load_work)
@@ -187,13 +191,13 @@ class MainWindow(QMainWindow):
 
     def view_operation_names(self) -> None:
         if self._check_show_names.isChecked():
-            self.is_show_names = True
+            self._show_names = True
         else:
-            self.is_show_names = False
+            self._show_names = False
 
         for operation in self._drag_operation_scenes:
-            operation.label.setOpacity(self.is_show_names)
-            operation.is_show_name = self.is_show_names
+            operation.label.setOpacity(self._show_names)
+            operation.show_name = self._show_names
 
     def _save_work(self) -> None:
         sfg = cast(SFG, self.sfg_widget.sfg)
@@ -265,7 +269,7 @@ class MainWindow(QMainWindow):
         self._load_sfg(sfg, positions)
         self._logger.info("Loaded SFG from path: " + str(module))
 
-    def _load_sfg(self, sfg, positions=None) -> None:
+    def _load_sfg(self, sfg: SFG, positions=None) -> None:
         if positions is None:
             positions = {}
 
@@ -276,29 +280,29 @@ class MainWindow(QMainWindow):
                 positions[op.graph_id][-1] if op.graph_id in positions else None,
             )
 
-        def connect_ports(ports):
+        def connect_ports(ports: Sequence[InputPort]):
             for port in ports:
                 for signal in port.signals:
-                    source = [
+                    sources = [
                         source
-                        for source in self._ports[
-                            self._drag_buttons[signal.source.operation]
-                        ]
+                        for source in self._drag_buttons[
+                            signal.source_operation
+                        ].port_list
                         if source.port is signal.source
                     ]
-                    destination = [
+                    destinations = [
                         destination
-                        for destination in self._ports[
-                            self._drag_buttons[signal.destination.operation]
-                        ]
+                        for destination in self._drag_buttons[
+                            signal.destination.operation
+                        ].port_list
                         if destination.port is signal.destination
                     ]
 
-                    if source and destination:
-                        self._connect_button(source[0], destination[0])
+                    if sources and destinations:
+                        self._connect_button(sources[0], destinations[0])
 
-            for port in self._pressed_ports:
-                port.select_port()
+            for pressed_port in self._pressed_ports:
+                pressed_port.select_port()
 
         for op in sfg.split():
             connect_ports(op.inputs)
@@ -312,33 +316,33 @@ class MainWindow(QMainWindow):
 
     def _create_recent_file_actions_and_menus(self):
         for i in range(self._max_recent_files):
-            recentFileAction = QAction(self._ui.recent_sfg)
-            recentFileAction.setVisible(False)
-            recentFileAction.triggered.connect(
-                lambda b=0, x=recentFileAction: self._open_recent_file(x)
+            recent_file_action = QAction(self._ui.recent_sfg)
+            recent_file_action.setVisible(False)
+            recent_file_action.triggered.connect(
+                lambda b=0, x=recent_file_action: self._open_recent_file(x)
             )
-            self._recent_files.append(recentFileAction)
-            self._ui.recent_sfg.addAction(recentFileAction)
+            self._recent_files_actions.append(recent_file_action)
+            self._ui.recent_sfg.addAction(recent_file_action)
 
         self._update_recent_file_list()
 
     def _update_recent_file_list(self):
         settings = QSettings()
 
-        rfp = settings.value("SFG/recentFiles")
+        rfp = cast(deque, settings.value("SFG/recentFiles"))
 
         # print(rfp)
         if rfp:
             dequelen = len(rfp)
             if dequelen > 0:
                 for i in range(dequelen):
-                    action = self._recent_files[i]
+                    action = self._recent_files_actions[i]
                     action.setText(rfp[i])
                     action.setData(QFileInfo(rfp[i]))
                     action.setVisible(True)
 
                 for i in range(dequelen, self._max_recent_files):
-                    self._recent_files[i].setVisible(False)
+                    self._recent_files_actions[i].setVisible(False)
 
     def _open_recent_file(self, action):
         self._load_from_file(action.data().filePath())
@@ -346,7 +350,7 @@ class MainWindow(QMainWindow):
     def _add_recent_file(self, module):
         settings = QSettings()
 
-        rfp = settings.value("SFG/recentFiles")
+        rfp = cast(deque, settings.value("SFG/recentFiles"))
 
         if rfp:
             if module not in rfp:
@@ -369,9 +373,8 @@ class MainWindow(QMainWindow):
         self._pressed_ports.clear()
         self._drag_buttons.clear()
         self._drag_operation_scenes.clear()
-        self._arrows.clear()
+        self._arrow_ports.clear()
         self._ports.clear()
-        self._signal_ports.clear()
         self._sfg_dict.clear()
         self._scene.clear()
         self._logger.info("Workspace cleared.")
@@ -379,11 +382,11 @@ class MainWindow(QMainWindow):
     def create_sfg_from_toolbar(self) -> None:
         inputs = []
         outputs = []
-        for op in self._pressed_operations:
-            if isinstance(op.operation, Input):
-                inputs.append(op.operation)
-            elif isinstance(op.operation, Output):
-                outputs.append(op.operation)
+        for pressed_op in self._pressed_operations:
+            if isinstance(pressed_op.operation, Input):
+                inputs.append(pressed_op.operation)
+            elif isinstance(pressed_op.operation, Output):
+                outputs.append(pressed_op.operation)
 
         name, accepted = QInputDialog.getText(
             self, "Create SFG", "Name: ", QLineEdit.Normal
@@ -401,100 +404,96 @@ class MainWindow(QMainWindow):
         self._logger.info("Created SFG with name: %s from selected operations." % name)
 
         def check_equality(signal: Signal, signal_2: Signal) -> bool:
-            source = cast(OutputPort, signal.source)
-            source2 = cast(OutputPort, signal_2.source)
-            dest = cast(InputPort, signal.destination)
-            dest2 = cast(InputPort, signal_2.destination)
+            source_operation = cast(Operation, signal.source_operation)
+            source_operation2 = cast(Operation, signal_2.source_operation)
+            dest_operation = cast(Operation, signal.destination_operation)
+            dest_operation2 = cast(Operation, signal_2.destination_operation)
             if not (
-                source.operation.type_name() == source2.operation.type_name()
-                and dest.operation.type_name() == dest2.operation.type_name()
+                source_operation.type_name() == source_operation2.type_name()
+                and dest_operation.type_name() == dest_operation2.type_name()
             ):
                 return False
 
             if (
-                hasattr(source.operation, "value")
-                and hasattr(source2.operation, "value")
-                and hasattr(dest.operation, "value")
-                and hasattr(dest2.operation, "value")
+                hasattr(source_operation, "value")
+                and hasattr(source_operation2, "value")
+                and hasattr(dest_operation, "value")
+                and hasattr(dest_operation2, "value")
             ):
                 if not (
-                    source.operation.value == source2.operation.value
-                    and dest.operation.value == dest2.operation.value
+                    source_operation.value == source_operation2.value
+                    and dest_operation.value == dest_operation2.value
                 ):
                     return False
 
             if (
-                hasattr(source.operation, "name")
-                and hasattr(source2.operation, "name")
-                and hasattr(dest.operation, "name")
-                and hasattr(dest2.operation, "name")
+                hasattr(source_operation, "name")
+                and hasattr(source_operation2, "name")
+                and hasattr(dest_operation, "name")
+                and hasattr(dest_operation2, "name")
             ):
                 if not (
-                    source.operation.name == source2.operation.name
-                    and dest.operation.name == dest2.operation.name
+                    source_operation.name == source_operation2.name
+                    and dest_operation.name == dest_operation2.name
                 ):
                     return False
 
             try:
-                _signal_source_index = [
-                    source.operation.outputs.index(port)
-                    for port in source.operation.outputs
+                signal_source_index = [
+                    source_operation.outputs.index(port)
+                    for port in source_operation.outputs
                     if signal in port.signals
                 ]
-                _signal_2_source_index = [
-                    source2.operation.outputs.index(port)
-                    for port in source2.operation.outputs
+                signal_2_source_index = [
+                    source_operation2.outputs.index(port)
+                    for port in source_operation2.outputs
                     if signal_2 in port.signals
                 ]
             except ValueError:
                 return False  # Signal output connections not matching
 
             try:
-                _signal_destination_index = [
-                    dest.operation.inputs.index(port)
-                    for port in dest.operation.inputs
+                signal_destination_index = [
+                    dest_operation.inputs.index(port)
+                    for port in dest_operation.inputs
                     if signal in port.signals
                 ]
-                _signal_2_destination_index = [
-                    dest2.operation.inputs.index(port)
-                    for port in dest2.operation.inputs
+                signal_2_destination_index = [
+                    dest_operation2.inputs.index(port)
+                    for port in dest_operation2.inputs
                     if signal_2 in port.signals
                 ]
             except ValueError:
                 return False  # Signal input connections not matching
 
             return (
-                _signal_source_index == _signal_2_source_index
-                and _signal_destination_index == _signal_2_destination_index
+                signal_source_index == signal_2_source_index
+                and signal_destination_index == signal_2_destination_index
             )
 
         for _pressed_op in self._pressed_operations:
             for operation in sfg.operations:
                 for input_ in operation.inputs:
                     for signal in input_.signals:
-                        for line in self._signal_ports:
-                            if check_equality(line.signal, signal):
-                                line.source.operation.operation = (
-                                    signal.source.operation
-                                )
-                                line.destination.operation.operation = (
-                                    signal.destination.operation
+                        for arrow in self._arrow_ports:
+                            if check_equality(arrow.signal, signal):
+                                arrow.set_source_operation(signal.source_operation)
+                                arrow.set_destination_operation(
+                                    signal.destination_operation
                                 )
 
                 for output_ in operation.outputs:
                     for signal in output_.signals:
-                        for line in self._signal_ports:
-                            if check_equality(line.signal, signal):
-                                line.source.operation.operation = (
-                                    signal.source.operation
-                                )
-                                line.destination.operation.operation = (
-                                    signal.destination.operation
+                        for arrow in self._arrow_ports:
+                            if check_equality(arrow.signal, signal):
+                                arrow.set_source_operation(signal.source_operation)
+                                arrow.set_destination_operation(
+                                    signal.destination_operation
                                 )
 
-        for op in self._pressed_operations:
-            op.setToolTip(sfg.name)
-            self._operation_to_sfg[op] = sfg
+        for pressed_op in self._pressed_operations:
+            pressed_op.setToolTip(sfg.name)
+            self._operation_to_sfg[pressed_op] = sfg
 
         self._sfg_dict[sfg.name] = sfg
 
@@ -541,6 +540,10 @@ class MainWindow(QMainWindow):
 
         self.add_operations_from_namespace(namespace, self._ui.custom_operations_list)
 
+    def _update(self):
+        self._scene.update()
+        self._graphics_view.update()
+
     def add_operation(
         self,
         op: Operation,
@@ -581,7 +584,7 @@ class MainWindow(QMainWindow):
                 "border-color: black; border-width: 2px"
             )
             attr_button.add_ports()
-            self._ports[attr_button] = attr_button.ports
+            self._ports[attr_button] = attr_button.port_list
 
             icon_path = os.path.join(
                 os.path.dirname(__file__),
@@ -608,7 +611,7 @@ class MainWindow(QMainWindow):
                 )
             attr_button_scene.setFlag(QGraphicsItem.ItemIsSelectable, True)
             operation_label = QGraphicsTextItem(op.name, attr_button_scene)
-            if not self.is_show_names:
+            if not self._show_names:
                 operation_label.setOpacity(0)
             operation_label.setTransformOriginPoint(
                 operation_label.boundingRect().center()
@@ -651,7 +654,7 @@ class MainWindow(QMainWindow):
         )
         self._logger.info("Finished refreshing operation list.")
 
-    def on_list_widget_item_clicked(self, item) -> None:
+    def _on_list_widget_item_clicked(self, item) -> None:
         self._create_operation_item(item)
 
     def keyPressEvent(self, event) -> None:
@@ -714,8 +717,8 @@ class MainWindow(QMainWindow):
         self._logger.info(
             "Connecting: %s -> %s."
             % (
-                source.operation.operation.type_name(),
-                destination.operation.operation.type_name(),
+                source.operation.type_name(),
+                destination.operation.type_name(),
             )
         )
         try:
@@ -723,18 +726,17 @@ class MainWindow(QMainWindow):
         except StopIteration:
             arrow = Arrow(source, destination, self)
 
-        if arrow not in self._signal_ports:
-            self._signal_ports[arrow] = []
+        if arrow not in self._arrow_ports:
+            self._arrow_ports[arrow] = []
 
-        self._signal_ports[arrow].append((source, destination))
+        self._arrow_ports[arrow].append((source, destination))
         self._scene.addItem(arrow)
-        self._arrows.append(arrow)
 
         self.update()
 
     def paintEvent(self, event) -> None:
-        for signal in self._signal_ports.keys():
-            signal.moveLine()
+        for arrow in self._arrow_ports:
+            arrow.update_arrow()
 
     def _select_operations(self) -> None:
         selected = [button.widget() for button in self._scene.selectedItems()]
@@ -796,7 +798,7 @@ class MainWindow(QMainWindow):
 
 def start_editor(sfg: Optional[SFG] = None):
     app = QApplication(sys.argv)
-    window = MainWindow()
+    window = SFGMainWindow()
     if sfg:
         window._load_sfg(sfg)
     window.show()
diff --git a/b_asic/GUI/port_button.py b/b_asic/GUI/port_button.py
index c743b3a7e9cb7a89d1f9d13b6a6d8eab6f936ec0..b41b301c872a22f77ccd94568de90c6ac17c0026 100644
--- a/b_asic/GUI/port_button.py
+++ b/b_asic/GUI/port_button.py
@@ -9,6 +9,7 @@ from qtpy.QtWidgets import QMenu, QPushButton
 
 if TYPE_CHECKING:
     from b_asic.GUI.drag_button import DragButton
+    from b_asic.operation import Operation
     from b_asic.port import Port
 
 
@@ -19,20 +20,19 @@ class PortButton(QPushButton):
     Parameters
     ----------
     name : str
-    operation : :class:`~b_asic.GUI.drag_button.DragButton`
+    operation_button : :class:`~b_asic.GUI.drag_button.DragButton`
     port : :class:`~b_asic.port.Port`
     """
 
     connectionRequested = Signal(QPushButton)
     moved = Signal()
 
-    def __init__(self, name: str, operation: "DragButton", port: "Port"):
-        super().__init__(name, parent=operation)
+    def __init__(self, name: str, operation_button: "DragButton", port: "Port"):
+        super().__init__(name, parent=operation_button)
         self.pressed = False
-        self._window = operation._window
+        self._window = operation_button._window
         self.port = port
-        self.operation = operation
-        self.clicked = 0
+        self._operation_button = operation_button
         self._m_drag = False
         self._m_press = False
         self.setAcceptDrops(True)
@@ -41,6 +41,11 @@ class PortButton(QPushButton):
         self.setStyleSheet("background-color: white")
         self.connectionRequested.connect(self._window._connect_callback)
 
+    @property
+    def operation(self) -> "Operation":
+        """Operation associated with PortButton."""
+        return self._operation_button.operation
+
     def contextMenuEvent(self, event):
         menu = QMenu()
         menu.addAction("Connect", lambda: self.connectionRequested.emit(self))
@@ -97,6 +102,19 @@ class PortButton(QPushButton):
         )
 
     def select_port(self, modifiers=None):
+        """
+        Select the port taking *modifiers* into account.
+
+        Parameters
+        ----------
+        modifiers : optional
+            Qt keyboard modifier.
+
+        Returns
+        -------
+
+
+        """
         if modifiers != Qt.KeyboardModifier.ControlModifier:
             for port in self._window._pressed_ports:
                 port._toggle_port(port.pressed)
@@ -111,5 +129,5 @@ class PortButton(QPushButton):
             else:
                 self._window._pressed_ports.append(self)
 
-        for signal in self._window._arrows:
-            signal.update()
+        for arrow in self._window._arrow_ports:
+            arrow.update()
diff --git a/b_asic/GUI/properties_window.py b/b_asic/GUI/properties_window.py
index b3707e8bbf2a5f9988b59de83b6ec7c00585f27b..cb6cda29b7c75303566c946b9a724e62aa0ab136 100644
--- a/b_asic/GUI/properties_window.py
+++ b/b_asic/GUI/properties_window.py
@@ -52,7 +52,7 @@ class PropertiesWindow(QDialog):
 
         self._show_name_layout = QHBoxLayout()
         self._check_show_name = QCheckBox("Show name?")
-        self._check_show_name.setChecked(self.operation.is_show_name)
+        self._check_show_name.setChecked(self.operation.show_name)
         self._check_show_name.setLayoutDirection(Qt.RightToLeft)
         self._check_show_name.setStyleSheet("spacing: 170px")
         self._show_name_layout.addWidget(self._check_show_name)
@@ -155,10 +155,10 @@ class PropertiesWindow(QDialog):
 
         if self._check_show_name.isChecked():
             self.operation.label.setOpacity(1)
-            self.operation.is_show_name = True
+            self.operation.show_name = True
         else:
             self.operation.label.setOpacity(0)
-            self.operation.is_show_name = False
+            self.operation.show_name = False
 
         self.operation.operation.set_latency_offsets(
             {
diff --git a/b_asic/GUI/select_sfg_window.py b/b_asic/GUI/select_sfg_window.py
index 69647b13bd7c4d9cf77416db4110f2a9a1ffc516..483e978b474169e77516f6b851a1c4164b1a29c1 100644
--- a/b_asic/GUI/select_sfg_window.py
+++ b/b_asic/GUI/select_sfg_window.py
@@ -7,13 +7,13 @@ from qtpy.QtCore import Qt, Signal
 from qtpy.QtWidgets import QComboBox, QDialog, QPushButton, QVBoxLayout
 
 if TYPE_CHECKING:
-    from b_asic.GUI.main_window import MainWindow
+    from b_asic.GUI.main_window import SFGMainWindow
 
 
 class SelectSFGWindow(QDialog):
     ok = Signal()
 
-    def __init__(self, window: "MainWindow"):
+    def __init__(self, window: "SFGMainWindow"):
         super().__init__()
         self._window = window
         self.setWindowFlags(Qt.WindowTitleHint | Qt.WindowCloseButtonHint)
diff --git a/b_asic/GUI/simulate_sfg_window.py b/b_asic/GUI/simulate_sfg_window.py
index 99b432208b610f4fc5fdb8fdd71a5b1607528c34..671e9fe690940f54044b3c632f564c822fcc7cde 100644
--- a/b_asic/GUI/simulate_sfg_window.py
+++ b/b_asic/GUI/simulate_sfg_window.py
@@ -1,7 +1,7 @@
 """
 B-ASIC window to simulate an SFG.
 """
-from typing import Dict
+from typing import TYPE_CHECKING, Dict
 
 from qtpy.QtCore import Qt, Signal
 from qtpy.QtWidgets import (
@@ -20,6 +20,9 @@ from qtpy.QtWidgets import (
 
 from b_asic.GUI.signal_generator_input import _GENERATOR_MAPPING
 
+if TYPE_CHECKING:
+    from b_asic.signal_flow_graph import SFG
+
 
 class SimulateSFGWindow(QDialog):
     simulate = Signal()
@@ -42,7 +45,15 @@ class SimulateSFGWindow(QDialog):
         self._input_grid = QGridLayout()
         self._input_files = {}
 
-    def add_sfg_to_dialog(self, sfg) -> None:
+    def add_sfg_to_dialog(self, sfg: "SFG") -> None:
+        """
+        Add a signal flow graph to the dialog.
+
+        Parameters
+        ----------
+        sfg : SFG
+            The signal flow graph to add.
+        """
         sfg_layout = QVBoxLayout()
         options_layout = QFormLayout()
 
@@ -105,6 +116,16 @@ class SimulateSFGWindow(QDialog):
         self._dialog_layout.addLayout(sfg_layout)
 
     def change_input_format(self, i: int, text: str) -> None:
+        """
+        Change the input format selector.
+
+        Parameters
+        ----------
+        i : int
+            Input number to change for.
+        text : str
+            Name of generator.
+        """
         grid = self._input_grid.itemAtPosition(i, 2)
         if grid:
             for j in reversed(range(grid.count())):
@@ -114,8 +135,6 @@ class SimulateSFGWindow(QDialog):
                     widget.hide()
             self._input_grid.removeItem(grid)
 
-        param_grid = QGridLayout()
-
         if text in _GENERATOR_MAPPING:
             param_grid = _GENERATOR_MAPPING[text](self._window._logger)
         else:
@@ -124,6 +143,7 @@ class SimulateSFGWindow(QDialog):
         self._input_grid.addLayout(param_grid, i, 2)
 
     def save_properties(self) -> None:
+        """Save the simulation properties and emit a signal to start the simulation."""
         for sfg, _properties in self._input_fields.items():
             ic_value = self._input_fields[sfg]["iteration_count"].value()
             if ic_value == 0:
@@ -159,4 +179,5 @@ class SimulateSFGWindow(QDialog):
 
     @property
     def properties(self) -> Dict:
+        """Return the simulation properties."""
         return self._properties
diff --git a/b_asic/GUI/simulation_worker.py b/b_asic/GUI/simulation_worker.py
index 455580f2b3f7d241b33ad56f6ece44f435fc35a5..73cfa308b1335e9aadfb15f303468473234ca903 100644
--- a/b_asic/GUI/simulation_worker.py
+++ b/b_asic/GUI/simulation_worker.py
@@ -24,6 +24,7 @@ class SimulationWorker(QObject):
         self._props = properties
 
     def start_simulation(self):
+        """Start simulation and emit signal when finished."""
         simulation = Simulation(self._sfg, input_providers=self._props["input_values"])
         simulation.run_for(
             self._props["iteration_count"],
diff --git a/b_asic/operation.py b/b_asic/operation.py
index dd5a93d0c3a8dd8818d50863e7169988b693943c..ccff1c5eb0551b321f83fff8da73f911a2e27d63 100644
--- a/b_asic/operation.py
+++ b/b_asic/operation.py
@@ -660,8 +660,8 @@ class AbstractOperation(Operation, AbstractGraphComponent):
             dict_ele = []
             for signal in inport.signals:
                 if signal.source:
-                    if signal.source.operation.graph_id:
-                        dict_ele.append(signal.source.operation.graph_id)
+                    if signal.source_operation.graph_id:
+                        dict_ele.append(signal.source_operation.graph_id)
                     else:
                         dict_ele.append(GraphID("no_id"))
                 else:
@@ -679,8 +679,8 @@ class AbstractOperation(Operation, AbstractGraphComponent):
             dict_ele = []
             for signal in outport.signals:
                 if signal.destination:
-                    if signal.destination.operation.graph_id:
-                        dict_ele.append(signal.destination.operation.graph_id)
+                    if signal.destination_operation.graph_id:
+                        dict_ele.append(signal.destination_operation.graph_id)
                     else:
                         dict_ele.append(GraphID("no_id"))
                 else:
@@ -893,7 +893,7 @@ class AbstractOperation(Operation, AbstractGraphComponent):
         Operations input ports.
         """
         return [
-            signal.source.operation for signal in self.input_signals if signal.source
+            signal.source_operation for signal in self.input_signals if signal.source
         ]
 
     @property
diff --git a/b_asic/schedule.py b/b_asic/schedule.py
index 1cd5678e26eb73d36bd36ce94e13f067413d4a7d..45da224bb92a627abd3c5b5de2dd6c5f97e31d9a 100644
--- a/b_asic/schedule.py
+++ b/b_asic/schedule.py
@@ -497,7 +497,7 @@ class Schedule:
                 tmp_prev_available = tmp_usage - new_slack
                 prev_available = tmp_prev_available % self._schedule_time
                 laps = new_slack // self._schedule_time
-                source_op = signal.source.operation
+                source_op = signal.source_operation
                 if new_usage < prev_available:
                     print("Incrementing input laps 1")
                     laps += 1
diff --git a/b_asic/scheduler_gui/__init__.py b/b_asic/scheduler_gui/__init__.py
index a66f3c344a08c410cb697209d4508b24a3ca2fad..fa401de50874f25c76ba2bd1dc2aa7d3bf7ca715 100644
--- a/b_asic/scheduler_gui/__init__.py
+++ b/b_asic/scheduler_gui/__init__.py
@@ -3,6 +3,6 @@ B-ASIC Scheduler-gui Module.
 
 Graphical user interface for B-ASIC scheduler.
 """
-from b_asic.scheduler_gui.main_window import MainWindow, start_scheduler
+from b_asic.scheduler_gui.main_window import start_scheduler
 
-__all__ = ['MainWindow', 'start_scheduler']
+__all__ = ['start_scheduler']
diff --git a/b_asic/scheduler_gui/main_window.py b/b_asic/scheduler_gui/main_window.py
index e02c7b8b017db368bfe7a53fdc7e5fb591f0bddb..98a2b2bdb8f2810b1f4410d06508c781539a7a5c 100644
--- a/b_asic/scheduler_gui/main_window.py
+++ b/b_asic/scheduler_gui/main_window.py
@@ -14,7 +14,7 @@ import webbrowser
 from collections import deque
 from copy import deepcopy
 from importlib.machinery import SourceFileLoader
-from typing import List, Optional, Union, cast
+from typing import Deque, List, Optional, cast
 
 # Qt/qtpy
 import qtpy
@@ -92,7 +92,7 @@ QCoreApplication.setApplicationName("B-ASIC Scheduling GUI")
 QCoreApplication.setApplicationVersion(__version__)
 
 
-class MainWindow(QMainWindow, Ui_MainWindow):
+class ScheduleMainWindow(QMainWindow, Ui_MainWindow):
     """Schedule of an SFG with scheduled Operations."""
 
     _scene: QGraphicsScene
@@ -150,9 +150,9 @@ class MainWindow(QMainWindow, Ui_MainWindow):
         self.splitter.setCollapsible(1, True)
 
         # Recent files
-        self.maxFileNr = 4
-        self.recentFilesList: List[QAction] = []
-        self.recentFilePaths = deque(maxlen=self.maxFileNr)
+        self._max_recent_files = 4
+        self._recent_files_actions: List[QAction] = []
+        self._recent_file_paths: Deque[str] = deque(maxlen=self._max_recent_files)
         self._create_recent_file_actions_and_menus()
 
     def _init_graphics(self) -> None:
@@ -650,34 +650,34 @@ class MainWindow(QMainWindow, Ui_MainWindow):
             log.error("'Operator' not found in info table. It may have been renamed.")
 
     def _create_recent_file_actions_and_menus(self):
-        for i in range(self.maxFileNr):
-            recentFileAction = QAction(self.menu_Recent_Schedule)
-            recentFileAction.setVisible(False)
-            recentFileAction.triggered.connect(
-                lambda b=0, x=recentFileAction: self._open_recent_file(x)
+        for i in range(self._max_recent_files):
+            recent_file_action = QAction(self.menu_Recent_Schedule)
+            recent_file_action.setVisible(False)
+            recent_file_action.triggered.connect(
+                lambda b=0, x=recent_file_action: self._open_recent_file(x)
             )
-            self.recentFilesList.append(recentFileAction)
-            self.menu_Recent_Schedule.addAction(recentFileAction)
+            self._recent_files_actions.append(recent_file_action)
+            self.menu_Recent_Schedule.addAction(recent_file_action)
 
         self._update_recent_file_list()
 
     def _update_recent_file_list(self):
         settings = QSettings()
 
-        rfp = settings.value("scheduler/recentFiles")
+        rfp = cast(deque, settings.value("scheduler/recentFiles"))
 
         # print(rfp)
         if rfp:
             dequelen = len(rfp)
             if dequelen > 0:
                 for i in range(dequelen):
-                    action = self.recentFilesList[i]
+                    action = self._recent_files_actions[i]
                     action.setText(rfp[i])
                     action.setData(QFileInfo(rfp[i]))
                     action.setVisible(True)
 
-                for i in range(dequelen, self.maxFileNr):
-                    self.recentFilesList[i].setVisible(False)
+                for i in range(dequelen, self._max_recent_files):
+                    self._recent_files_actions[i].setVisible(False)
 
     def _open_recent_file(self, action):
         self._load_from_file(action.data().filePath())
@@ -685,13 +685,13 @@ class MainWindow(QMainWindow, Ui_MainWindow):
     def _add_recent_file(self, module):
         settings = QSettings()
 
-        rfp = settings.value("scheduler/recentFiles")
+        rfp = cast(deque, settings.value("scheduler/recentFiles"))
 
         if rfp:
             if module not in rfp:
                 rfp.append(module)
         else:
-            rfp = deque(maxlen=self.maxFileNr)
+            rfp = deque(maxlen=self._max_recent_files)
             rfp.append(module)
 
         settings.setValue("scheduler/recentFiles", rfp)
@@ -701,7 +701,7 @@ class MainWindow(QMainWindow, Ui_MainWindow):
 
 def start_scheduler(schedule: Optional[Schedule] = None) -> None:
     app = QApplication(sys.argv)
-    window = MainWindow()
+    window = ScheduleMainWindow()
     if schedule:
         window.open(schedule)
     window.show()
diff --git a/b_asic/signal.py b/b_asic/signal.py
index 5f4e0948282096eae5f43659ebdf223e5578e9f1..86b2325e7061e2c00e8332ec2142425383042695 100644
--- a/b_asic/signal.py
+++ b/b_asic/signal.py
@@ -77,6 +77,20 @@ class Signal(AbstractGraphComponent):
         """Return the destination InputPort of the signal."""
         return self._destination
 
+    @property
+    def source_operation(self) -> Optional["Operation"]:
+        """Return the source Operation of the signal."""
+        if self._source is not None:
+            return self._source.operation
+        return None
+
+    @property
+    def destination_operation(self) -> Optional["Operation"]:
+        """Return the destination Operation of the signal."""
+        if self._destination is not None:
+            return self._destination.operation
+        return None
+
     def set_source(self, source: Union["OutputPort", "Signal", "Operation"]) -> None:
         """
         Connect to the entered source OutputPort.
diff --git a/b_asic/utils.py b/b_asic/utils.py
index 99a5f5169a60f67d518ba2db722e73bd6ebe25ac..baa3b29afded1d29eb60271cf80815722d7b4892 100644
--- a/b_asic/utils.py
+++ b/b_asic/utils.py
@@ -9,7 +9,7 @@ def interleave(*args) -> List[Num]:
     """
     Interleave a number of arrays.
 
-    For the input ``interleave([1, 2], [3, 4])``, return ``[1, 2, 3, 4]``.
+    For the input ``interleave([1, 2], [3, 4])``, return ``[1, 3, 2, 4]``.
 
     Parameters
     ----------
diff --git a/test/test_gui.py b/test/test_gui.py
index beb22644ce4681be795e8c1018178cdc1bd581b4..95c39693c245d785260206207254ddb90f3055dc 100644
--- a/test/test_gui.py
+++ b/test/test_gui.py
@@ -3,7 +3,7 @@ from qtpy import QtCore
 from qtpy.QtWidgets import QInputDialog
 
 try:
-    import b_asic.GUI as GUI
+    from b_asic.GUI.main_window import SFGMainWindow
 except ImportError:
     pytestmark = pytest.mark.skip("Qt not setup")
 
@@ -12,14 +12,14 @@ from b_asic.special_operations import Input, Output
 
 
 def test_start(qtbot):
-    widget = GUI.MainWindow()
+    widget = SFGMainWindow()
     qtbot.addWidget(widget)
 
     widget.exit_app()
 
 
 def test_load(qtbot, datadir):
-    widget = GUI.MainWindow()
+    widget = SFGMainWindow()
     qtbot.addWidget(widget)
     widget._load_from_file(datadir.join('twotapfir.py'))
     assert 'twotapfir' in widget._sfg_dict
@@ -31,7 +31,7 @@ def test_load(qtbot, datadir):
 
 
 def test_flip(qtbot, datadir):
-    widget = GUI.MainWindow()
+    widget = SFGMainWindow()
     qtbot.addWidget(widget)
     widget._load_from_file(datadir.join('twotapfir.py'))
     sfg = widget._sfg_dict['twotapfir']
@@ -44,7 +44,7 @@ def test_flip(qtbot, datadir):
 
 
 def test_sfg_invalidated_by_remove_of_operation(qtbot, datadir):
-    widget = GUI.MainWindow()
+    widget = SFGMainWindow()
     qtbot.addWidget(widget)
     widget._load_from_file(datadir.join('twotapfir.py'))
     sfg = widget._sfg_dict['twotapfir']
@@ -59,7 +59,7 @@ def test_sfg_invalidated_by_remove_of_operation(qtbot, datadir):
 
 
 def test_sfg_invalidated_by_deleting_of_operation(qtbot, datadir):
-    widget = GUI.MainWindow()
+    widget = SFGMainWindow()
     qtbot.addWidget(widget)
     widget._load_from_file(datadir.join('twotapfir.py'))
     sfg = widget._sfg_dict['twotapfir']
@@ -76,7 +76,7 @@ def test_sfg_invalidated_by_deleting_of_operation(qtbot, datadir):
 
 
 def test_select_operation(qtbot, datadir):
-    widget = GUI.MainWindow()
+    widget = SFGMainWindow()
     qtbot.addWidget(widget)
     widget._load_from_file(datadir.join('twotapfir.py'))
     sfg = widget._sfg_dict['twotapfir']
@@ -134,7 +134,7 @@ def test_select_operation(qtbot, datadir):
 def test_help_dialogs(qtbot):
     # Smoke test to open up the "help dialogs"
     # Should really test doing this through the menus an/or closing them
-    widget = GUI.MainWindow()
+    widget = SFGMainWindow()
     qtbot.addWidget(widget)
 
     widget.display_faq_page()
@@ -151,7 +151,7 @@ def test_help_dialogs(qtbot):
 def test_simulate(qtbot, datadir):
     # Smoke test to open up the "Simulate SFG" and run default simulation
     # Should really test all different tests
-    widget = GUI.MainWindow()
+    widget = SFGMainWindow()
     qtbot.addWidget(widget)
     widget._load_from_file(datadir.join('twotapfir.py'))
     assert 'twotapfir' in widget._sfg_dict
@@ -167,7 +167,7 @@ def test_simulate(qtbot, datadir):
 def test_properties_window_smoke_test(qtbot, datadir):
     # Smoke test to open up the _properties window
     # Should really check that the contents are correct and changes works etc
-    widget = GUI.MainWindow()
+    widget = SFGMainWindow()
     qtbot.addWidget(widget)
     widget._load_from_file(datadir.join('twotapfir.py'))
     sfg = widget._sfg_dict['twotapfir']
@@ -184,7 +184,7 @@ def test_properties_window_smoke_test(qtbot, datadir):
 def test_properties_window_change_name(qtbot, datadir):
     # Smoke test to open up the _properties window
     # Should really check that the contents are correct and changes works etc
-    widget = GUI.MainWindow()
+    widget = SFGMainWindow()
     qtbot.addWidget(widget)
     widget._load_from_file(datadir.join('twotapfir.py'))
     sfg = widget._sfg_dict['twotapfir']
@@ -206,7 +206,7 @@ def test_properties_window_change_name(qtbot, datadir):
 
 
 def test_add_operation_and_create_sfg(qtbot, monkeypatch):
-    widget = GUI.MainWindow()
+    widget = SFGMainWindow()
     qtbot.addWidget(widget)
     in1 = Input()
     sqrt = SquareRoot()
@@ -221,10 +221,10 @@ def test_add_operation_and_create_sfg(qtbot, monkeypatch):
     for op in (in1, sqrt, out1):
         assert op in widget._drag_buttons
     # No signals
-    assert not widget._arrows
+    assert not widget._arrow_ports
 
     # Click on first port
-    in1_port = widget._ports[widget._drag_buttons[in1]][0]
+    in1_port = widget._drag_buttons[in1].port_list[0]
     qtbot.mouseClick(
         in1_port,
         QtCore.Qt.MouseButton.LeftButton,
@@ -232,7 +232,7 @@ def test_add_operation_and_create_sfg(qtbot, monkeypatch):
     assert len(widget._pressed_ports) == 1
 
     # Click on second port
-    sqrt_in_port = widget._ports[widget._drag_buttons[sqrt]][0]
+    sqrt_in_port = widget._drag_buttons[sqrt].port_list[0]
     qtbot.mouseClick(
         sqrt_in_port,
         QtCore.Qt.MouseButton.LeftButton,
@@ -245,16 +245,16 @@ def test_add_operation_and_create_sfg(qtbot, monkeypatch):
     # Not sure why this won't work
     # qtbot.keyClick(widget, QtCore.Qt.Key.Key_Space, delay=10)
     # Still one selected!?
-    assert len(widget._arrows) == 1
+    assert len(widget._arrow_ports) == 1
 
     # Click on first port
-    sqrt_out_port = widget._ports[widget._drag_buttons[sqrt]][1]
+    sqrt_out_port = widget._drag_buttons[sqrt].port_list[1]
     qtbot.mouseClick(
         sqrt_out_port,
         QtCore.Qt.MouseButton.LeftButton,
     )
     # Click on second port
-    out1_port = widget._ports[widget._drag_buttons[out1]][0]
+    out1_port = widget._drag_buttons[out1].port_list[0]
     qtbot.mouseClick(
         out1_port,
         QtCore.Qt.MouseButton.LeftButton,
@@ -262,7 +262,7 @@ def test_add_operation_and_create_sfg(qtbot, monkeypatch):
     )
     # Connect
     widget._connect_callback()
-    assert len(widget._arrows) == 2
+    assert len(widget._arrow_ports) == 2
 
     # Select input op
     qtbot.mouseClick(
diff --git a/test/test_scheduler_gui.py b/test/test_scheduler_gui.py
index 9dc4abcecc490c911105b94185e9a39565091cf7..d6ca2517aed686edd846f5be7ae40ece18b36415 100644
--- a/test/test_scheduler_gui.py
+++ b/test/test_scheduler_gui.py
@@ -4,13 +4,13 @@ from b_asic.core_operations import Addition, ConstantMultiplication
 from b_asic.schedule import Schedule
 
 try:
-    import b_asic.scheduler_gui as GUI
+    from b_asic.scheduler_gui.main_window import ScheduleMainWindow
 except ImportError:
     pytestmark = pytest.mark.skip("Qt not setup")
 
 
 def test_start(qtbot):
-    widget = GUI.MainWindow()
+    widget = ScheduleMainWindow()
     qtbot.addWidget(widget)
 
     widget.exit_app()
@@ -18,11 +18,9 @@ def test_start(qtbot):
 
 def test_load_schedule(qtbot, sfg_simple_filter):
     sfg_simple_filter.set_latency_of_type(Addition.type_name(), 5)
-    sfg_simple_filter.set_latency_of_type(
-        ConstantMultiplication.type_name(), 4
-    )
+    sfg_simple_filter.set_latency_of_type(ConstantMultiplication.type_name(), 4)
 
-    widget = GUI.MainWindow()
+    widget = ScheduleMainWindow()
     qtbot.addWidget(widget)
     schedule = Schedule(sfg_simple_filter)
     widget.open(schedule)