Change type aliases for support in Python 3.8

Farama-Foundation · Apr 25, 2024 · cfbb24d · cfbb24d
1 parent 3f1f2b8
commit cfbb24d
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 7 deletions.
diff --git a/metaworld/envs/mujoco/env_dict.py b/metaworld/envs/mujoco/env_dict.py
@@ -10,14 +10,17 @@
 from typing import Sequence, Union
 
 import numpy as np
+from typing_extensions import TypeAlias
 
 from metaworld.envs.mujoco.sawyer_xyz import SawyerXYZEnv, v2
 
 # Utils
 
-EnvDict = Typing_OrderedDict[str, type[SawyerXYZEnv]]
-TrainTestEnvDict = Typing_OrderedDict[Literal["train", "test"], EnvDict]
-EnvArgsKwargsDict = Dict[str, Dict[Literal["args", "kwargs"], Union[List, Dict]]]
+EnvDict: TypeAlias = "Typing_OrderedDict[str, type[SawyerXYZEnv]]"
+TrainTestEnvDict: TypeAlias = "Typing_OrderedDict[Literal['train', 'test'], EnvDict]"
+EnvArgsKwargsDict: TypeAlias = (
+    "Dict[str, Dict[Literal['args', 'kwargs'], Union[List, Dict]]]"
+)
 
 ENV_CLS_MAP = {
     "assembly-v2": v2.SawyerNutAssemblyEnvV2,

diff --git a/metaworld/envs/mujoco/sawyer_xyz/sawyer_xyz_env.py b/metaworld/envs/mujoco/sawyer_xyz/sawyer_xyz_env.py
@@ -13,11 +13,12 @@
 from gymnasium.spaces import Box, Discrete, Space
 from gymnasium.utils import seeding
 from gymnasium.utils.ezpickle import EzPickle
+from typing_extensions import TypeAlias
 
 from metaworld.envs.mujoco.utils import reward_utils
 from metaworld.types import XYZ, EnvironmentStateDict, ObservationDict, Task
 
-RenderMode = Literal["human", "rgb_array", "depth_array"]
+RenderMode: TypeAlias = "Literal['human', 'rgb_array', 'depth_array']"
 
 
 class SawyerMocapBase(mjenv_gym):

diff --git a/metaworld/envs/mujoco/utils/reward_utils.py b/metaworld/envs/mujoco/utils/reward_utils.py
@@ -21,7 +21,7 @@
     "tanh_squared",
 ]
 
-X = TypeVar("X", float, npt.NDArray[Any], np.floating[Any])
+X = TypeVar("X", float, npt.NDArray, np.floating)
 
 
 def _sigmoids(x: X, value_at_1: float, sigmoid: SIGMOID_TYPE) -> X:

diff --git a/metaworld/types.py b/metaworld/types.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 import numpy.typing as npt
-from typing_extensions import NotRequired, TypedDict
+from typing_extensions import NotRequired, TypeAlias, TypedDict
 
 
 class Task(NamedTuple):
@@ -17,7 +17,7 @@ class Task(NamedTuple):
     data: bytes  # Contains env parameters like random_init and *a* goal
 
 
-XYZ = Tuple[float, float, float]
+XYZ: TypeAlias = "Tuple[float, float, float]"
 """A 3D coordinate."""