compatibility with python3.9 (#55)

* compatibility with python 3.9

* bump version

Author: Howuhh

src/xminigrid/__init__.py

@@ -2,7 +2,7 @@
 from .registration import make, register, registered_environments
 
 # TODO: add __all__
-__version__ = "0.9.2"
+__version__ = "0.9.3"
 
 # ---------- XLand-MiniGrid environments ----------
 

src/xminigrid/benchmarks.py

@@ -4,7 +4,7 @@
 import os
 import pickle
 import urllib.request
-from typing import Callable
+from typing import Callable, Union
 
 import jax
 import jax.numpy as jnp
@@ -38,7 +38,7 @@ class Benchmark(struct.PyTreeNode):
     def num_rulesets(self) -> int:
         return len(self.goals)
 
-    def get_ruleset(self, ruleset_id: int | jax.Array) -> RuleSet:
+    def get_ruleset(self, ruleset_id: Union[int, jax.Array]) -> RuleSet:
         return get_ruleset(self.goals, self.rules, self.init_tiles, ruleset_id)
 
     def sample_ruleset(self, key: jax.Array) -> RuleSet:
@@ -114,7 +114,7 @@ def get_ruleset(
     goals: jax.Array,
     rules: jax.Array,
     init_tiles: jax.Array,
-    ruleset_id: int | jax.Array,
+    ruleset_id: Union[int, jax.Array],
 ) -> RuleSet:
     goal = jax.lax.dynamic_index_in_dim(goals, ruleset_id, keepdims=False)
     rules = jax.lax.dynamic_index_in_dim(rules, ruleset_id, keepdims=False)

src/xminigrid/core/goals.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+from typing import Union
 
 import jax
 import jax.numpy as jnp
@@ -14,7 +15,7 @@
 
 
 def check_goal(
-    encoding: jax.Array, grid: GridState, agent: AgentState, action: int | jax.Array, position: jax.Array
+    encoding: jax.Array, grid: GridState, agent: AgentState, action: Union[int, jax.Array], position: jax.Array
 ) -> jax.Array:
     check = jax.lax.switch(
         encoding[0],
@@ -45,7 +46,7 @@ def check_goal(
 class BaseGoal(struct.PyTreeNode):
     @abc.abstractmethod
     def __call__(
-        self, grid: GridState, agent: AgentState, action: int | jax.Array, position: jax.Array
+        self, grid: GridState, agent: AgentState, action: Union[int, jax.Array], position: jax.Array
     ) -> jax.Array: ...
 
     @classmethod

src/xminigrid/core/grid.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Callable
+from typing import Callable, Optional, Union
 
 import jax
 import jax.numpy as jnp
@@ -158,7 +158,7 @@ def coordinates_mask(grid: GridState, address: tuple[IntOrArray, IntOrArray], co
     return mask
 
 
-def sample_coordinates(key: jax.Array, grid: GridState, num: int, mask: jax.Array | None = None) -> jax.Array:
+def sample_coordinates(key: jax.Array, grid: GridState, num: int, mask: Optional[jax.Array] = None) -> jax.Array:
     if mask is None:
         mask = jnp.ones((grid.shape[0], grid.shape[1]), dtype=jnp.bool_)
 

src/xminigrid/core/rules.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+from typing import Union
 
 import jax
 import jax.numpy as jnp
@@ -18,7 +19,7 @@
 # In general, we need a way to select specific function/class based on ID number.
 # We can not just decode without evaluation, as then return type will be different between branches
 def check_rule(
-    encodings: jax.Array, grid: GridState, agent: AgentState, action: int | jax.Array, position: jax.Array
+    encodings: jax.Array, grid: GridState, agent: AgentState, action: Union[int, jax.Array], position: jax.Array
 ) -> tuple[GridState, AgentState]:
     def _check(carry, encoding):
         grid, agent = carry
@@ -51,7 +52,7 @@ def _check(carry, encoding):
 class BaseRule(struct.PyTreeNode):
     @abc.abstractmethod
     def __call__(
-        self, grid: GridState, agent: AgentState, action: int | jax.Array, position: jax.Array
+        self, grid: GridState, agent: AgentState, action: Union[int, jax.Array], position: jax.Array
     ) -> tuple[GridState, AgentState]: ...
 
     @classmethod

src/xminigrid/environment.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import abc
-from typing import Any, Generic, Optional, TypeVar
+from typing import Any, Generic, Optional, TypeVar, Union
 
 import jax
 import jax.numpy as jnp
@@ -40,7 +40,7 @@ def default_params(self, **kwargs: Any) -> EnvParamsT: ...
     def num_actions(self, params: EnvParamsT) -> int:
         return int(NUM_ACTIONS)
 
-    def observation_shape(self, params: EnvParamsT) -> tuple[int, int, int] | dict[str, Any]:
+    def observation_shape(self, params: EnvParamsT) -> Union[tuple[int, int, int], dict[str, Any]]:
         return params.view_size, params.view_size, NUM_LAYERS
 
     @abc.abstractmethod
@@ -89,7 +89,7 @@ def step(self, params: EnvParamsT, timestep: TimeStep[EnvCarryT], action: IntOrA
         )
         return timestep
 
-    def render(self, params: EnvParamsT, timestep: TimeStep[EnvCarryT]) -> np.ndarray | str:
+    def render(self, params: EnvParamsT, timestep: TimeStep[EnvCarryT]) -> Union[np.ndarray, str]:
         if params.render_mode == "rgb_array":
             return rgb_render(np.asarray(timestep.state.grid), timestep.state.agent, params.view_size)
         elif params.render_mode == "rich_text":

src/xminigrid/manual_control.py

@@ -2,6 +2,7 @@
 
 import argparse
 import os
+from typing import Optional
 
 import imageio.v3 as iio
 import jax
@@ -24,7 +25,7 @@ def __init__(
         env_params: EnvParamsT,
         agent_view: bool = False,
         save_video: bool = False,
-        video_path: str | None = None,
+        video_path: Optional[str] = None,
         video_format: str = ".mp4",
         video_fps: int = 8,
     ):

src/xminigrid/rendering/rgb_render.py

@@ -1,6 +1,7 @@
 # TODO: this is rendering mostly ported or adapted from the original Minigrid. A bit dirty right now...
 import functools
 import math
+from typing import Optional, Union
 
 import numpy as np
 
@@ -181,7 +182,7 @@ def _render_player(img: np.ndarray, direction: int):
 
 
 # TODO: add highlight for can_see_through_walls=Fasle
-def get_highlight_mask(grid: np.ndarray, agent: AgentState | None, view_size: int) -> np.ndarray:
+def get_highlight_mask(grid: np.ndarray, agent: Optional[AgentState], view_size: int) -> np.ndarray:
     mask = np.zeros((grid.shape[0] + 2 * view_size, grid.shape[1] + 2 * view_size), dtype=np.bool_)
     if agent is None:
         return mask
@@ -207,7 +208,7 @@ def get_highlight_mask(grid: np.ndarray, agent: AgentState | None, view_size: in
 
 @functools.cache
 def render_tile(
-    tile: tuple, agent_direction: int | None = None, highlight: bool = False, tile_size: int = 32, subdivs: int = 3
+    tile: tuple, agent_direction: Optional[int] = None, highlight: bool = False, tile_size: int = 32, subdivs: int = 3
 ) -> np.ndarray:
     img = np.full((tile_size * subdivs, tile_size * subdivs, 3), dtype=np.uint8, fill_value=255)
     # draw tile
@@ -228,7 +229,7 @@ def render_tile(
 # WARN: will NOT work under jit and needed for debugging/presentation mainly.
 def render(
     grid: np.ndarray,
-    agent: AgentState | None = None,
+    agent: Optional[AgentState] = None,
     view_size: IntOrArray = 7,
     tile_size: IntOrArray = 32,
 ) -> np.ndarray:

src/xminigrid/rendering/text_render.py

@@ -1,3 +1,5 @@
+from typing import Optional, Union
+
 import jax
 import jax.numpy as jnp
 
@@ -55,7 +57,7 @@ def _wrap_with_color(string: str, color: str) -> str:
 
 
 # WARN: will NOT work under jit and needed for debugging mainly.
-def render(grid: jax.Array, agent: AgentState | None = None) -> str:
+def render(grid: jax.Array, agent: Optional[AgentState] = None) -> str:
     string = ""
 
     for y in range(grid.shape[0]):

src/xminigrid/rendering/utils.py

@@ -2,12 +2,12 @@
 from __future__ import annotations
 
 import math
-from typing import Callable
+from typing import Callable, Union
 
 import numpy as np
 from typing_extensions import TypeAlias
 
-Color: TypeAlias = tuple[int, int, int] | int | np.ndarray
+Color: TypeAlias = Union[tuple[int, int, int], int, np.ndarray]
 Point: TypeAlias = tuple[float, float]  # | np.ndarray