Fix tests

mixtape/core/analysis/ray_utils/wrappers.py

@@ -1,7 +1,14 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 import numpy as np
 from ray.rllib.env import PettingZooEnv
 from ray.rllib.env.wrappers.pettingzoo_env import ParallelPettingZooEnv
 
+if TYPE_CHECKING:
+    import numpy.typing as npt
+
 
 class PZWrapper(PettingZooEnv):
     """Temporary wrapper for PettingZooEnv.
@@ -11,12 +18,13 @@ class PZWrapper(PettingZooEnv):
     render function, but PettingZoo environments no longer accept an argument.
     """
 
-    def render(self) -> np._typing.NDArray[np.uint8]:
+    # Superclass return type is incorrect:
+    # https://github.com/Farama-Foundation/Gymnasium/issues/845#issuecomment-2535827312
+    def render(self) -> npt.NDArray[np.uint8]:  # type: ignore[override]
         """RGB image given the current observation.
 
         Returns:
-            np._typing.NDArray[np.uint8]: A numpy uint8 3D array (image) to
-                                          render.
+            A numpy uint8 3D array (image) to render.
         """
         return self.env.render()
 
@@ -29,11 +37,12 @@ class ParallelPZWrapper(ParallelPettingZooEnv):
     render function, but PettingZoo environments no longer accept an argument.
     """
 
-    def render(self) -> np._typing.NDArray[np.uint8]:
+    # Superclass return type is incorrect:
+    # https://github.com/Farama-Foundation/Gymnasium/issues/845#issuecomment-2535827312
+    def render(self) -> npt.NDArray[np.uint8]:  # type: ignore[override]
         """RGB image given the current observation.
 
         Returns:
-            np._typing.NDArray[np.uint8]: A numpy uint8 3D array (image) to
-                                          render.
+            A numpy uint8 3D array (image) to render.
         """
         return self.par_env.render()

mixtape/core/ray_utils/constants.py

@@ -1,3 +1,5 @@
+from typing import Literal
+
 from django.db import models