From e759b97f8663ee6ed431f958c05b09d02e1523f4 Mon Sep 17 00:00:00 2001
From: TJ <tbruno25@gmail.com>
Date: Fri, 19 Jul 2024 10:05:00 -0700
Subject: [PATCH] Fix mypy

---
 src/can_explorer/controllers.py | 7 ++++---
 src/can_explorer/models.py      | 3 ++-
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/can_explorer/controllers.py b/src/can_explorer/controllers.py
index b6183ca..c8b88ce 100644
--- a/src/can_explorer/controllers.py
+++ b/src/can_explorer/controllers.py
@@ -3,6 +3,7 @@
 import enum
 
 import can
+from can.bus import BusABC
 
 from can_explorer.configs import Default
 from can_explorer.models import PlotModel
@@ -19,7 +20,7 @@ def __init__(
         self,
         model: PlotModel,
         view: MainView,
-        bus: can.bus.BusABC | None = None,
+        bus: BusABC | None = None,
         refresh_rate: float | None = Default.REFRESH_RATE,
     ) -> None:
         self.model = model
@@ -35,7 +36,7 @@ def state(self) -> State:
         return self._state
 
     @property
-    def bus(self) -> can.bus.BusABC | None:
+    def bus(self) -> BusABC:
         if self._bus is None:
             raise RuntimeError("Must apply settings before starting")
         return self._bus
@@ -43,7 +44,7 @@ def bus(self) -> can.bus.BusABC | None:
     def is_active(self) -> bool:
         return bool(self.state)
 
-    def set_bus(self, bus: can.BusABC) -> None:
+    def set_bus(self, bus: BusABC) -> None:
         """
         Set CAN bus to use during controller loop.
         """
diff --git a/src/can_explorer/models.py b/src/can_explorer/models.py
index 4bc9313..9903b8a 100644
--- a/src/can_explorer/models.py
+++ b/src/can_explorer/models.py
@@ -1,4 +1,5 @@
 from collections import defaultdict, deque
+from typing import DefaultDict
 
 import can
 
@@ -22,7 +23,7 @@ def __getitem__(self, index) -> tuple:  # type: ignore [override]
 
 class PlotModel:
     def __init__(self) -> None:
-        self._data = defaultdict(PayloadBuffer)
+        self._data: DefaultDict[int, PayloadBuffer] = defaultdict(PayloadBuffer)
         self._len = Default.BUFFER_SIZE
 
     def add_message(self, message: can.Message) -> None: