From 7f8fa241dd501dbab7dba5c7e95f66702ecf039b Mon Sep 17 00:00:00 2001
From: Jiayi Weng <trinkle23897@gmail.com>
Date: Sat, 25 Mar 2023 22:01:09 -0700
Subject: [PATCH] making pettingzoo a core dep instead of optional req (#837)

close #831
---
 .github/workflows/extra_sys.yml | 2 +-
 .github/workflows/gputest.yml   | 4 ++--
 .github/workflows/pytest.yml    | 2 +-
 setup.py                        | 2 +-
 tianshou/__init__.py            | 2 +-
 tianshou/env/__init__.py        | 6 +-----
 tianshou/env/gym_wrappers.py    | 4 ++--
 tianshou/env/utils.py           | 5 +----
 tianshou/env/venvs.py           | 6 +-----
 9 files changed, 11 insertions(+), 22 deletions(-)

diff --git a/.github/workflows/extra_sys.yml b/.github/workflows/extra_sys.yml
index db945cb..3dca090 100644
--- a/.github/workflows/extra_sys.yml
+++ b/.github/workflows/extra_sys.yml
@@ -9,7 +9,7 @@ jobs:
     strategy:
       matrix:
         os: [macos-latest, windows-latest]
-        python-version: [3.7, 3.8]
+        python-version: [3.8, 3.9]
     steps:
       - name: Cancel previous run
         uses: styfle/cancel-workflow-action@0.11.0
diff --git a/.github/workflows/gputest.yml b/.github/workflows/gputest.yml
index 18ff687..9bb9469 100644
--- a/.github/workflows/gputest.yml
+++ b/.github/workflows/gputest.yml
@@ -12,10 +12,10 @@ jobs:
         with:
           access_token: ${{ github.token }}
       - uses: actions/checkout@v3
-      - name: Set up Python 3.8
+      - name: Set up Python 3.10
         uses: actions/setup-python@v4
         with:
-          python-version: 3.8
+          python-version: "3.10"
       - name: Upgrade pip
         run: |
           python -m pip install --upgrade pip setuptools wheel
diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index c08d8ec..1dc128e 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -8,7 +8,7 @@ jobs:
     if: "!contains(github.event.head_commit.message, 'ci skip')"
     strategy:
       matrix:
-        python-version: [3.7, 3.8, 3.9]
+        python-version: ["3.8", "3.9", "3.10"]
     steps:
       - name: Cancel previous run
         uses: styfle/cancel-workflow-action@0.11.0
diff --git a/setup.py b/setup.py
index 253be20..8ae0050 100644
--- a/setup.py
+++ b/setup.py
@@ -23,6 +23,7 @@ def get_install_requires() -> str:
         "numba>=0.51.0",
         "h5py>=2.10.0",  # to match tensorflow's minimal requirements
         "packaging",
+        "pettingzoo>=1.22",
     ]
 
 
@@ -47,7 +48,6 @@ def get_extras_require() -> str:
             "doc8",
             "scipy",
             "pillow",
-            "pettingzoo>=1.22",
             "pygame>=2.1.0",  # pettingzoo test cases pistonball
             "pymunk>=6.2.1",  # pettingzoo test cases pistonball
             "nni>=2.3,<3.0",  # expect breaking changes at next major version
diff --git a/tianshou/__init__.py b/tianshou/__init__.py
index e039d6e..e3f08b1 100644
--- a/tianshou/__init__.py
+++ b/tianshou/__init__.py
@@ -1,6 +1,6 @@
 from tianshou import data, env, exploration, policy, trainer, utils
 
-__version__ = "0.5.0"
+__version__ = "0.5.1"
 
 __all__ = [
     "env",
diff --git a/tianshou/env/__init__.py b/tianshou/env/__init__.py
index a00c3cd..049ccf4 100644
--- a/tianshou/env/__init__.py
+++ b/tianshou/env/__init__.py
@@ -5,6 +5,7 @@ from tianshou.env.gym_wrappers import (
     MultiDiscreteToDiscrete,
     TruncatedAsTerminated,
 )
+from tianshou.env.pettingzoo_env import PettingZooEnv
 from tianshou.env.venv_wrappers import VectorEnvNormObs, VectorEnvWrapper
 from tianshou.env.venvs import (
     BaseVectorEnv,
@@ -14,11 +15,6 @@ from tianshou.env.venvs import (
     SubprocVectorEnv,
 )
 
-try:
-    from tianshou.env.pettingzoo_env import PettingZooEnv
-except ImportError:
-    pass
-
 __all__ = [
     "BaseVectorEnv",
     "DummyVectorEnv",
diff --git a/tianshou/env/gym_wrappers.py b/tianshou/env/gym_wrappers.py
index c9ce66a..f3aec8f 100644
--- a/tianshou/env/gym_wrappers.py
+++ b/tianshou/env/gym_wrappers.py
@@ -26,7 +26,7 @@ class ContinuousToDiscrete(gym.ActionWrapper):
             dtype=object
         )
 
-    def action(self, act: np.ndarray) -> np.ndarray:  # type: ignore
+    def action(self, act: np.ndarray) -> np.ndarray:
         # modify act
         assert len(act.shape) <= 2, f"Unknown action format with shape {act.shape}."
         if len(act.shape) == 1:
@@ -50,7 +50,7 @@ class MultiDiscreteToDiscrete(gym.ActionWrapper):
             self.bases[i] = self.bases[i - 1] * nvec[-i]
         self.action_space = gym.spaces.Discrete(np.prod(nvec))
 
-    def action(self, act: np.ndarray) -> np.ndarray:  # type: ignore
+    def action(self, act: np.ndarray) -> np.ndarray:
         converted_act = []
         for b in np.flip(self.bases):
             converted_act.append(act // b)
diff --git a/tianshou/env/utils.py b/tianshou/env/utils.py
index cbd36d9..ac67101 100644
--- a/tianshou/env/utils.py
+++ b/tianshou/env/utils.py
@@ -4,10 +4,7 @@ import cloudpickle
 import gymnasium
 import numpy as np
 
-try:
-    from tianshou.env.pettingzoo_env import PettingZooEnv
-except ImportError:
-    PettingZooEnv = None  # type: ignore
+from tianshou.env.pettingzoo_env import PettingZooEnv
 
 if TYPE_CHECKING:
     import gym
diff --git a/tianshou/env/venvs.py b/tianshou/env/venvs.py
index cc39ba3..09b2090 100644
--- a/tianshou/env/venvs.py
+++ b/tianshou/env/venvs.py
@@ -5,6 +5,7 @@ import gymnasium as gym
 import numpy as np
 import packaging
 
+from tianshou.env.pettingzoo_env import PettingZooEnv
 from tianshou.env.utils import ENV_TYPE, gym_new_venv_step_type
 from tianshou.env.worker import (
     DummyEnvWorker,
@@ -13,11 +14,6 @@ from tianshou.env.worker import (
     SubprocEnvWorker,
 )
 
-try:
-    from tianshou.env.pettingzoo_env import PettingZooEnv
-except ImportError:
-    PettingZooEnv = None  # type: ignore
-
 try:
     import gym as old_gym