Fix type annotations of callable to Callable (openai#259)

gymnasium/envs/registration.py

@@ -45,7 +45,7 @@
 )
 
 
-def load(name: str) -> callable:
+def load(name: str) -> Callable:
     """Loads an environment with name and returns an environment creation function.
 
     Args:

gymnasium/utils/play.py

@@ -321,7 +321,7 @@ class PlayPlot:
     """
 
     def __init__(
-        self, callback: callable, horizon_timesteps: int, plot_names: List[str]
+        self, callback: Callable, horizon_timesteps: int, plot_names: List[str]
     ):
         """Constructor of :class:`PlayPlot`.
 

gymnasium/vector/__init__.py

@@ -1,7 +1,8 @@
 """Module for vector environments."""
-from typing import Iterable, List, Optional, Union
+from typing import Callable, Iterable, List, Optional, Union
 
 import gymnasium as gym
+from gymnasium.core import Env
 from gymnasium.vector.async_vector_env import AsyncVectorEnv
 from gymnasium.vector.sync_vector_env import SyncVectorEnv
 from gymnasium.vector.vector_env import VectorEnv, VectorEnvWrapper
@@ -14,7 +15,7 @@ def make(
     id: str,
     num_envs: int = 1,
     asynchronous: bool = True,
-    wrappers: Optional[Union[callable, List[callable]]] = None,
+    wrappers: Optional[Union[Callable[[Env], Env], List[Callable[[Env], Env]]]] = None,
     disable_env_checker: Optional[bool] = None,
     **kwargs,
 ) -> VectorEnv:
@@ -43,12 +44,12 @@ def make(
         The vectorized environment.
     """
 
-    def create_env(env_num: int):
+    def create_env(env_num: int) -> Callable[[], Env]:
         """Creates an environment that can enable or disable the environment checker."""
         # If the env_num > 0 then disable the environment checker otherwise use the parameter
         _disable_env_checker = True if env_num > 0 else disable_env_checker
 
-        def _make_env():
+        def _make_env() -> Env:
             env = gym.envs.registration.make(
                 id,
                 disable_env_checker=_disable_env_checker,

gymnasium/vector/async_vector_env.py

@@ -4,13 +4,13 @@
 import time
 from copy import deepcopy
 from enum import Enum
-from typing import List, Optional, Sequence, Tuple, Union
+from typing import Callable, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 
 import gymnasium as gym
 from gymnasium import logger
-from gymnasium.core import ObsType
+from gymnasium.core import Env, ObsType
 from gymnasium.error import (
     AlreadyPendingCallError,
     ClosedEnvironmentError,
@@ -59,14 +59,14 @@ class AsyncVectorEnv(VectorEnv):
 
     def __init__(
         self,
-        env_fns: Sequence[callable],
+        env_fns: Sequence[Callable[[], Env]],
         observation_space: Optional[gym.Space] = None,
         action_space: Optional[gym.Space] = None,
         shared_memory: bool = True,
         copy: bool = True,
         context: Optional[str] = None,
         daemon: bool = True,
-        worker: Optional[callable] = None,
+        worker: Optional[Callable] = None,
     ):
         """Vectorized environment that runs multiple environments in parallel.
 

gymnasium/vector/utils/misc.py

@@ -1,6 +1,11 @@
 """Miscellaneous utilities."""
+from __future__ import annotations
+
 import contextlib
 import os
+from collections.abc import Callable
+
+from gymnasium.core import Env
 
 
 __all__ = ["CloudpickleWrapper", "clear_mpi_env_vars"]
@@ -9,7 +14,7 @@
 class CloudpickleWrapper:
     """Wrapper that uses cloudpickle to pickle and unpickle the result."""
 
-    def __init__(self, fn: callable):
+    def __init__(self, fn: Callable[[], Env]):
         """Cloudpickle wrapper for a function."""
         self.fn = fn
 

gymnasium/vector/utils/numpy_utils.py

@@ -1,7 +1,7 @@
 """Numpy utility functions: concatenate space samples and create empty array."""
 from collections import OrderedDict
 from functools import singledispatch
-from typing import Iterable, Union
+from typing import Callable, Iterable, Union
 
 import numpy as np
 
@@ -84,7 +84,7 @@ def _concatenate_custom(space, items, out):
 
 @singledispatch
 def create_empty_array(
-    space: Space, n: int = 1, fn: callable = np.zeros
+    space: Space, n: int = 1, fn: Callable[..., np.ndarray] = np.zeros
 ) -> Union[tuple, dict, np.ndarray]:
     """Create an empty (possibly nested) numpy array.
 

tests/spaces/test_spaces.py

@@ -3,7 +3,7 @@
 import json  # note: ujson fails this test due to float equality
 import pickle
 import tempfile
-from typing import List, Union
+from typing import Callable, List, Union
 
 import numpy as np
 import pytest
@@ -283,7 +283,7 @@ def test_space_sample_mask(space: Space, mask, n_trials: int = 100):
     elif isinstance(space, MultiDiscrete):
         # Due to the multi-axis capability of MultiDiscrete, these functions need to be recursive and that the expected / observed numpy are of non-regular shapes
         def _generate_frequency(
-            _dim: Union[np.ndarray, int], _mask, func: callable
+            _dim: Union[np.ndarray, int], _mask, func: Callable
         ) -> List:
             if isinstance(_dim, np.ndarray):
                 return [

tests/testing_env.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import types
+from collections.abc import Callable
 from typing import Any
 
 import gymnasium as gym
@@ -45,9 +46,9 @@ def __init__(
         self,
         action_space: spaces.Space = spaces.Box(0, 1, (1,)),
         observation_space: spaces.Space = spaces.Box(0, 1, (1,)),
-        reset_func: callable = basic_reset_func,
-        step_func: callable = new_step_func,
-        render_func: callable = basic_render_func,
+        reset_func: Callable = basic_reset_func,
+        step_func: Callable = new_step_func,
+        render_func: Callable = basic_render_func,
         metadata: dict[str, Any] = {"render_modes": []},
         render_mode: str | None = None,
         spec: EnvSpec = EnvSpec(

tests/utils/test_env_checker.py

@@ -1,7 +1,7 @@
 """Tests that the `env_checker` runs as expects and all errors are possible."""
 import re
 import warnings
-from typing import Tuple, Union
+from typing import Callable, Tuple, Union
 
 import numpy as np
 import pytest
@@ -106,7 +106,7 @@ def _reset_default_seed(self: GenericTestEnv, seed="Error", options=None):
         ],
     ],
 )
-def test_check_reset_seed(test, func: callable, message: str):
+def test_check_reset_seed(test, func: Callable, message: str):
     """Tests the check reset seed function works as expected."""
     if test is UserWarning:
         with pytest.warns(
@@ -175,7 +175,7 @@ def _return_info_not_dict(self, seed=None, options=None):
         ],
     ],
 )
-def test_check_reset_return_type(test, func: callable, message: str):
+def test_check_reset_return_type(test, func: Callable, message: str):
     """Tests the check `env.reset()` function has a correct return type."""
 
     with pytest.raises(test, match=f"^{re.escape(message)}$"):
@@ -194,7 +194,7 @@ def test_check_reset_return_type(test, func: callable, message: str):
         ],
     ],
 )
-def test_check_reset_return_info_deprecation(test, func: callable, message: str):
+def test_check_reset_return_info_deprecation(test, func: Callable, message: str):
     """Tests that return_info has been correct deprecated as an argument to `env.reset()`."""
 
     with pytest.warns(test, match=f"^\\x1b\\[33mWARN: {re.escape(message)}\\x1b\\[0m$"):

tests/utils/test_passive_env_checker.py

@@ -1,6 +1,6 @@
 import re
 import warnings
-from typing import Dict, Union
+from typing import Callable, Dict, Union
 
 import numpy as np
 import pytest
@@ -297,7 +297,7 @@ def _reset_result(self, seed=None, options=None):
         ],
     ],
 )
-def test_passive_env_reset_checker(test, func: callable, message: str, kwargs: Dict):
+def test_passive_env_reset_checker(test, func: Callable, message: str, kwargs: Dict):
     """Tests the passive env reset check"""
     if test is UserWarning:
         with pytest.warns(
@@ -376,7 +376,7 @@ def _modified_step(
     ],
 )
 def test_passive_env_step_checker(
-    test: Union[UserWarning, type], func: callable, message: str
+    test: Union[UserWarning, type], func: Callable, message: str
 ):
     """Tests the passive env step checker."""
     if test is UserWarning: