From 279d45accfed92d11046fb3a127a55137002d9a0 Mon Sep 17 00:00:00 2001
From: Oscar Gustafsson <oscar.gustafsson@gmail.com>
Date: Fri, 19 Apr 2024 15:19:18 +0200
Subject: [PATCH] Improve typing

---
 b_asic/GUI/drag_button.py           |  6 ++++--
 b_asic/scheduler_gui/main_window.py | 14 +++++++++++++-
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/b_asic/GUI/drag_button.py b/b_asic/GUI/drag_button.py
index cb6a9ca0..36d03c07 100644
--- a/b_asic/GUI/drag_button.py
+++ b/b_asic/GUI/drag_button.py
@@ -19,6 +19,8 @@ from b_asic.operation import Operation
 from b_asic.port import InputPort
 
 if TYPE_CHECKING:
+    from qtpy.QtWidgets import QGraphicsTextItem
+
     from b_asic.GUI.main_window import SFGMainWindow
 
 
@@ -90,13 +92,13 @@ class DragButton(QPushButton):
         """Return the type name of the underlying operation."""
         return self.operation.type_name()
 
-    def add_label(self, label: str) -> None:
+    def add_label(self, label: "QGraphicsTextItem") -> None:
         """
         Add label to button.
 
         Parameters
         ----------
-        label : src
+        label : QGraphicsTextItem
             The label to add.
         """
         self.label = label
diff --git a/b_asic/scheduler_gui/main_window.py b/b_asic/scheduler_gui/main_window.py
index b331abcc..7a0c4114 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 defaultdict, deque
 from copy import deepcopy
 from importlib.machinery import SourceFileLoader
-from typing import TYPE_CHECKING, Deque, List, Optional, cast
+from typing import TYPE_CHECKING, Deque, List, Optional, cast, overload
 
 # Qt/qtpy
 import qtpy
@@ -1006,6 +1006,18 @@ class ScheduleMainWindow(QMainWindow, Ui_MainWindow):
             self.actionToggle_full_screen.setIcon(get_icon('full-screen-exit'))
 
 
+@overload
+def start_scheduler(schedule: Schedule) -> Schedule: ...
+
+
+@overload
+def start_scheduler(schedule: None) -> Optional[Schedule]: ...
+
+
+@overload
+def start_scheduler() -> Optional[Schedule]: ...
+
+
 def start_scheduler(schedule: Optional[Schedule] = None) -> Optional[Schedule]:
     """
     Start scheduler GUI.
-- 
GitLab