diff --git a/mixtape/core/analysis/ray_utils/wrappers.py b/mixtape/core/analysis/ray_utils/wrappers.py index 7213153..fa28e0d 100644 --- a/mixtape/core/analysis/ray_utils/wrappers.py +++ b/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() diff --git a/mixtape/core/ray_utils/constants.py b/mixtape/core/ray_utils/constants.py index f93ee08..2b0bc47 100644 --- a/mixtape/core/ray_utils/constants.py +++ b/mixtape/core/ray_utils/constants.py @@ -1,3 +1,5 @@ +from typing import Literal + from django.db import models