Remove redundant 'if TYPE_CHECKING' statements

src/primaite/agents/agent_abc.py

@@ -4,8 +4,9 @@ from __future__ import annotations
 import json
 from abc import ABC, abstractmethod
 from datetime import datetime
+from logging import Logger
 from pathlib import Path
-from typing import Any, Dict, Optional, TYPE_CHECKING, Union
+from typing import Any, Dict, Optional, Union
 from uuid import uuid4
 
 import primaite
@@ -16,10 +17,7 @@ from primaite.data_viz.session_plots import plot_av_reward_per_episode
 from primaite.environment.primaite_env import Primaite
 from primaite.utils.session_metadata_parser import parse_session_metadata
 
-if TYPE_CHECKING:
-    from logging import Logger
-
-_LOGGER: "Logger" = getLogger(__name__)
+_LOGGER: Logger = getLogger(__name__)
 
 
 def get_session_path(session_timestamp: datetime) -> Path:

src/primaite/agents/rllib.py

@@ -4,8 +4,9 @@ from __future__ import annotations
 import json
 import shutil
 from datetime import datetime
+from logging import Logger
 from pathlib import Path
-from typing import Any, Callable, Dict, Optional, TYPE_CHECKING, Union
+from typing import Any, Callable, Dict, Optional, Union
 from uuid import uuid4
 
 from ray.rllib.algorithms import Algorithm
@@ -19,10 +20,7 @@ from primaite.agents.agent_abc import AgentSessionABC
 from primaite.common.enums import AgentFramework, AgentIdentifier
 from primaite.environment.primaite_env import Primaite
 
-if TYPE_CHECKING:
-    from logging import Logger
-
-_LOGGER: "Logger" = getLogger(__name__)
+_LOGGER: Logger = getLogger(__name__)
 
 
 # TODO: verify type of env_config

src/primaite/agents/sb3.py

@@ -2,8 +2,9 @@
 from __future__ import annotations
 
 import json
+from logging import Logger
 from pathlib import Path
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 from stable_baselines3 import A2C, PPO
@@ -14,10 +15,7 @@ from primaite.agents.agent_abc import AgentSessionABC
 from primaite.common.enums import AgentFramework, AgentIdentifier
 from primaite.environment.primaite_env import Primaite
 
-if TYPE_CHECKING:
-    from logging import Logger
-
-_LOGGER: "Logger" = getLogger(__name__)
+_LOGGER: Logger = getLogger(__name__)
 
 
 class SB3Agent(AgentSessionABC):

src/primaite/agents/simple.py

@@ -1,12 +1,10 @@
 # Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
-from typing import TYPE_CHECKING
+
+import numpy as np
 
 from primaite.agents.hardcoded_abc import HardCodedAgentSessionABC
 from primaite.agents.utils import get_new_action, transform_action_acl_enum, transform_action_node_enum
 
-if TYPE_CHECKING:
-    import numpy as np
-
 
 class RandomAgent(HardCodedAgentSessionABC):
     """
@@ -15,7 +13,7 @@ class RandomAgent(HardCodedAgentSessionABC):
     Get a completely random action from the action space.
     """
 
-    def _calculate_action(self, obs: "np.ndarray") -> int:
+    def _calculate_action(self, obs: np.ndarray) -> int:
         return self._env.action_space.sample()
 
 
@@ -26,7 +24,7 @@ class DummyAgent(HardCodedAgentSessionABC):
     All action spaces setup so dummy action is always 0 regardless of action type used.
     """
 
-    def _calculate_action(self, obs: "np.ndarray") -> int:
+    def _calculate_action(self, obs: np.ndarray) -> int:
         return 0
 
 
@@ -37,7 +35,7 @@ class DoNothingACLAgent(HardCodedAgentSessionABC):
     A valid ACL action that has no effect; does nothing.
     """
 
-    def _calculate_action(self, obs: "np.ndarray") -> int:
+    def _calculate_action(self, obs: np.ndarray) -> int:
         nothing_action = ["NONE", "ALLOW", "ANY", "ANY", "ANY", "ANY"]
         nothing_action = transform_action_acl_enum(nothing_action)
         nothing_action = get_new_action(nothing_action, self._env.action_dict)
@@ -52,7 +50,7 @@ class DoNothingNodeAgent(HardCodedAgentSessionABC):
     A valid Node action that has no effect; does nothing.
     """
 
-    def _calculate_action(self, obs: "np.ndarray") -> int:
+    def _calculate_action(self, obs: np.ndarray) -> int:
         nothing_action = [1, "NONE", "ON", 0]
         nothing_action = transform_action_node_enum(nothing_action)
         nothing_action = get_new_action(nothing_action, self._env.action_dict)