Continue Adding Typehints

src/primaite/common/custom_typing.py

@@ -1,8 +1,8 @@
-from typing import Type, Union
+from typing import TypeVar
 
 from primaite.nodes.active_node import ActiveNode
 from primaite.nodes.passive_node import PassiveNode
 from primaite.nodes.service_node import ServiceNode
 
-NodeUnion: Type = Union[ActiveNode, PassiveNode, ServiceNode]
+NodeUnion = TypeVar("NodeUnion", ServiceNode, ActiveNode, PassiveNode)
 """A Union of ActiveNode, PassiveNode, and ServiceNode."""

src/primaite/common/protocol.py

@@ -5,17 +5,17 @@
 class Protocol(object):
     """Protocol class."""
 
-    def __init__(self, _name):
+    def __init__(self, _name: str) -> None:
         """
         Initialise a protocol.
 
         :param _name: The name of the protocol
         :type _name: str
         """
-        self.name = _name
-        self.load = 0  # bps
+        self.name: str = _name
+        self.load: int = 0  # bps
 
-    def get_name(self):
+    def get_name(self) -> str:
         """
         Gets the protocol name.
 
@@ -24,7 +24,7 @@ class Protocol(object):
         """
         return self.name
 
-    def get_load(self):
+    def get_load(self) -> int:
         """
         Gets the protocol load.
 
@@ -33,7 +33,7 @@ class Protocol(object):
         """
         return self.load
 
-    def add_load(self, _load):
+    def add_load(self, _load: int) -> None:
         """
         Adds load to the protocol.
 
@@ -42,6 +42,6 @@ class Protocol(object):
         """
         self.load += _load
 
-    def clear_load(self):
+    def clear_load(self) -> None:
         """Clears the load on this protocol."""
         self.load = 0

src/primaite/common/service.py

@@ -15,12 +15,12 @@ class Service(object):
         :param port: The service port.
         :param software_state: The service SoftwareState.
         """
-        self.name = name
-        self.port = port
-        self.software_state = software_state
-        self.patching_count = 0
+        self.name: str = name
+        self.port: str = port
+        self.software_state: SoftwareState = software_state
+        self.patching_count: int = 0
 
-    def reduce_patching_count(self):
+    def reduce_patching_count(self) -> None:
         """Reduces the patching count for the service."""
         self.patching_count -= 1
         if self.patching_count <= 0: