Repository: sotetsuk/pgx

Issue:
`import pgx` is slow
- visualization code loads all environments
- Chess codes run some time-consuming preparation 


Discussion:
Where is the bottle neck?

- `from pgx._test import api_test`
- `from pgx._visualizer import save_svg`
- `from pgx._visualizer import save_svg_animation`
- `from pgx._visualizer import set_visualization_config,`
- `from pgx.core import Env, EnvId, State, available_games, make`

-> visualizer
例えば、囲碁の可視化をするだけでも、全部のゲームの実装がロードされているのが問題。

## 解決方法

- Stateだけ別ファイルにまとめる（多くの実装はロードされない）
- 可視化の条件分岐を違う方法で行う

Resolving patch:
diff --git a/pgx/_visualizer.py b/pgx/_visualizer.py
--- a/pgx/_visualizer.py
+++ b/pgx/_visualizer.py
@@ -19,20 +19,6 @@
 
 import svgwrite  # type: ignore
 
-from ._dwg.animalshogi import AnimalShogiState, _make_animalshogi_dwg
-from ._dwg.backgammon import BackgammonState, _make_backgammon_dwg
-from ._dwg.bridge_bidding import BridgeBiddingState, _make_bridge_dwg
-from ._dwg.chess import ChessState, _make_chess_dwg
-from ._dwg.connect_four import ConnectFourState, _make_connect_four_dwg
-from ._dwg.go import GoState, _make_go_dwg
-from ._dwg.hex import HexState, _make_hex_dwg
-from ._dwg.kuhn_poker import KuhnPokerState, _make_kuhnpoker_dwg
-from ._dwg.leduc_holdem import LeducHoldemState, _make_leducHoldem_dwg
-from ._dwg.othello import OthelloState, _make_othello_dwg
-from ._dwg.play2048 import Play2048State, _make_2048_dwg
-from ._dwg.shogi import ShogiState, _make_shogi_dwg
-from ._dwg.sparrow_mahjong import SparrowMahjongState, _make_sparrowmahjong_dwg
-from ._dwg.tictactoe import TictactoeState, _make_tictactoe_dwg
 from .core import State
 
 ColorTheme = Literal["light", "dark"]
@@ -198,12 +184,14 @@ def get_dwg(
         dwg.add(group)
         return dwg
 
-    def _set_config_by_state(self, _state):  # noqa: C901
-        if isinstance(_state, AnimalShogiState):
+    def _set_config_by_state(self, _state: State):  # noqa: C901
+        if _state.env_id == "animal_shogi":
+            from ._dwg.animalshogi import _make_animalshogi_dwg
+
             self.config["GRID_SIZE"] = 60
             self.config["BOARD_WIDTH"] = 4
             self.config["BOARD_HEIGHT"] = 4
-            self._make_dwg_group = _make_animalshogi_dwg
+            self._make_dwg_group = _make_animalshogi_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -227,11 +215,13 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "",
                 )
-        elif isinstance(_state, BackgammonState):
+        elif _state.env_id == "backgammon":
+            from ._dwg.backgammon import _make_backgammon_dwg
+
             self.config["GRID_SIZE"] = 25
             self.config["BOARD_WIDTH"] = 17
             self.config["BOARD_HEIGHT"] = 14
-            self._make_dwg_group = _make_backgammon_dwg
+            self._make_dwg_group = _make_backgammon_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -255,11 +245,13 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "gray",
                 )
-        elif isinstance(_state, BridgeBiddingState):
+        elif _state.env_id == "bridge_bidding":
+            from ._dwg.bridge_bidding import _make_bridge_dwg
+
             self.config["GRID_SIZE"] = 50
             self.config["BOARD_WIDTH"] = 14
             self.config["BOARD_HEIGHT"] = 10
-            self._make_dwg_group = _make_bridge_dwg
+            self._make_dwg_group = _make_bridge_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -283,11 +275,13 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "black",
                 )
-        elif isinstance(_state, ChessState):
+        elif _state.env_id == "chess":
+            from ._dwg.chess import _make_chess_dwg
+
             self.config["GRID_SIZE"] = 50
             self.config["BOARD_WIDTH"] = 8
             self.config["BOARD_HEIGHT"] = 8
-            self._make_dwg_group = _make_chess_dwg
+            self._make_dwg_group = _make_chess_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -311,11 +305,13 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "",
                 )
-        elif isinstance(_state, ConnectFourState):
+        elif _state.env_id == "connect_four":
+            from ._dwg.connect_four import _make_connect_four_dwg
+
             self.config["GRID_SIZE"] = 35
             self.config["BOARD_WIDTH"] = 7
             self.config["BOARD_HEIGHT"] = 7
-            self._make_dwg_group = _make_connect_four_dwg
+            self._make_dwg_group = _make_connect_four_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -339,15 +335,19 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "gray",
                 )
-        elif isinstance(_state, GoState):
+        elif _state.env_id in ("go-9x9", "go-19x19"):
+            from ._dwg.go import _make_go_dwg
+
             self.config["GRID_SIZE"] = 25
             try:
-                self.config["BOARD_WIDTH"] = int(_state.size[0])
-                self.config["BOARD_HEIGHT"] = int(_state.size[0])
+                self.config["BOARD_WIDTH"] = int(_state.size[0])  # type:ignore
+                self.config["BOARD_HEIGHT"] = int(
+                    _state.size[0]  # type:ignore
+                )
             except IndexError:
-                self.config["BOARD_WIDTH"] = int(_state.size)
-                self.config["BOARD_HEIGHT"] = int(_state.size)
-            self._make_dwg_group = _make_go_dwg
+                self.config["BOARD_WIDTH"] = int(_state.size)  # type:ignore
+                self.config["BOARD_HEIGHT"] = int(_state.size)  # type:ignore
+            self._make_dwg_group = _make_go_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -371,23 +371,27 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "",
                 )
-        elif isinstance(_state, HexState):
+        elif _state.env_id == "hex":
             import jax.numpy as jnp
 
-            from pgx._dwg.hex import four_dig
+            from ._dwg.hex import _make_hex_dwg, four_dig
 
             self.config["GRID_SIZE"] = 30
             try:
-                self.config["BOARD_WIDTH"] = four_dig(_state.size[0] * 1.5)
+                self.config["BOARD_WIDTH"] = four_dig(
+                    _state.size[0] * 1.5  # type:ignore
+                )
                 self.config["BOARD_HEIGHT"] = four_dig(
-                    _state.size[0] * jnp.sqrt(3) / 2
+                    _state.size[0] * jnp.sqrt(3) / 2  # type:ignore
                 )
             except IndexError:
-                self.config["BOARD_WIDTH"] = four_dig(_state.size * 1.5)
+                self.config["BOARD_WIDTH"] = four_dig(
+                    _state.size * 1.5  # type:ignore
+                )
                 self.config["BOARD_HEIGHT"] = four_dig(
-                    _state.size * jnp.sqrt(3) / 2
+                    _state.size * jnp.sqrt(3) / 2  # type:ignore
                 )
-            self._make_dwg_group = _make_hex_dwg
+            self._make_dwg_group = _make_hex_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -411,11 +415,13 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "lightgray",
                 )
-        elif isinstance(_state, KuhnPokerState):
+        elif _state.env_id == "kuhn_poker":
+            from ._dwg.kuhn_poker import _make_kuhnpoker_dwg
+
             self.config["GRID_SIZE"] = 30
             self.config["BOARD_WIDTH"] = 8
             self.config["BOARD_HEIGHT"] = 8
-            self._make_dwg_group = _make_kuhnpoker_dwg
+            self._make_dwg_group = _make_kuhnpoker_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -439,11 +445,13 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "",
                 )
-        elif isinstance(_state, LeducHoldemState):
+        elif _state.env_id == "leduc_holdem":
+            from ._dwg.leduc_holdem import _make_leducHoldem_dwg
+
             self.config["GRID_SIZE"] = 30
             self.config["BOARD_WIDTH"] = 8
             self.config["BOARD_HEIGHT"] = 8
-            self._make_dwg_group = _make_leducHoldem_dwg
+            self._make_dwg_group = _make_leducHoldem_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -467,11 +475,13 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "",
                 )
-        elif isinstance(_state, OthelloState):
+        elif _state.env_id == "othello":
+            from ._dwg.othello import _make_othello_dwg
+
             self.config["GRID_SIZE"] = 30
             self.config["BOARD_WIDTH"] = 8
             self.config["BOARD_HEIGHT"] = 8
-            self._make_dwg_group = _make_othello_dwg
+            self._make_dwg_group = _make_othello_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -495,11 +505,13 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "",
                 )
-        elif isinstance(_state, Play2048State):
+        elif _state.env_id == "2048":
+            from ._dwg.play2048 import _make_2048_dwg
+
             self.config["GRID_SIZE"] = 50
             self.config["BOARD_WIDTH"] = 4
             self.config["BOARD_HEIGHT"] = 4
-            self._make_dwg_group = _make_2048_dwg
+            self._make_dwg_group = _make_2048_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -523,11 +535,13 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "black",
                 )
-        elif isinstance(_state, ShogiState):
+        elif _state.env_id == "shogi":
+            from ._dwg.shogi import _make_shogi_dwg
+
             self.config["GRID_SIZE"] = 50
             self.config["BOARD_WIDTH"] = 10
             self.config["BOARD_HEIGHT"] = 9
-            self._make_dwg_group = _make_shogi_dwg
+            self._make_dwg_group = _make_shogi_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -545,11 +559,13 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "black",
                     "",
                 )
-        elif isinstance(_state, SparrowMahjongState):
+        elif _state.env_id == "sparrow_mahjong":
+            from ._dwg.sparrow_mahjong import _make_sparrowmahjong_dwg
+
             self.config["GRID_SIZE"] = 50
             self.config["BOARD_WIDTH"] = 15
             self.config["BOARD_HEIGHT"] = 10
-            self._make_dwg_group = _make_sparrowmahjong_dwg
+            self._make_dwg_group = _make_sparrowmahjong_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -573,11 +589,13 @@ def _set_config_by_state(self, _state):  # noqa: C901
                     "silver",
                     "black",
                 )
-        elif isinstance(_state, TictactoeState):
+        elif _state.env_id == "tic_tac_toe":
+            from ._dwg.tictactoe import _make_tictactoe_dwg
+
             self.config["GRID_SIZE"] = 60
             self.config["BOARD_WIDTH"] = 3
             self.config["BOARD_HEIGHT"] = 3
-            self._make_dwg_group = _make_tictactoe_dwg
+            self._make_dwg_group = _make_tictactoe_dwg  # type:ignore
             if (
                 self.config["COLOR_THEME"] is None
                 and self.config["COLOR_THEME"] == "dark"
@@ -597,93 +615,124 @@ def _set_config_by_state(self, _state):  # noqa: C901
         else:
             assert False
 
-    def _get_nth_state(self, _states, _i):
-        if isinstance(_states, AnimalShogiState):
+    # TODO: simplify me
+    def _get_nth_state(self, _states: State, _i):
+        if _states.env_id == "animal_shogi":
+            from ._dwg.animalshogi import AnimalShogiState
+
             return AnimalShogiState(
                 turn=_states.turn[_i],  # type:ignore
-                board=_states.board[_i],
-                hand=_states.hand[_i],
+                board=_states.board[_i],  # type:ignore
+                hand=_states.hand[_i],  # type:ignore
             )
-        elif isinstance(_states, BackgammonState):
+        elif _states.env_id == "backgammon":
+            from ._dwg.backgammon import BackgammonState
+
             return BackgammonState(
                 turn=_states.turn[_i],  # type:ignore
-                board=_states.board[_i],
+                board=_states.board[_i],  # type:ignore
             )
-        elif isinstance(_states, ConnectFourState):
+        elif _states.env_id == "connect_four":
+            from ._dwg.connect_four import ConnectFourState
+
             return ConnectFourState(  # type:ignore
-                turn=_states.turn[_i],
-                board=_states.board[_i],
+                turn=_states.turn[_i],  # type:ignore
+                board=_states.board[_i],  # type:ignore
             )
-        elif isinstance(_states, ChessState):
+        elif _states.env_id == "chess":
+            from ._dwg.chess import ChessState
+
             return ChessState(
                 turn=_states.turn[_i],  # type:ignore
-                board=_states.board[_i],
+                board=_states.board[_i],  # type:ignore
             )
-        elif isinstance(_states, BridgeBiddingState):
+        elif _states.env_id == "bridge_bidding":
+            from ._dwg.bridge_bidding import BridgeBiddingState
+
             return BridgeBiddingState(  # type:ignore
-                turn=_states.turn[_i],
-                dealer=_states.dealer[_i],
-                current_player=_states.current_player[_i],
-                hand=_states.hand[_i],
-                bidding_history=_states.bidding_history[_i],
-                vul_NS=_states.vul_NS[_i],
-                vul_EW=_states.vul_EW[_i],
+                turn=_states.turn[_i],  # type:ignore
+                dealer=_states.dealer[_i],  # type:ignore
+                current_player=_states.current_player[_i],  # type:ignore
+                hand=_states.hand[_i],  # type:ignore
+                bidding_history=_states.bidding_history[_i],  # type:ignore
+                vul_NS=_states.vul_NS[_i],  # type:ignore
+                vul_EW=_states.vul_EW[_i],  # type:ignore
             )
-        elif isinstance(_states, GoState):
+        elif _states.env_id in ("go-9x9", "go-19x19"):
+            from ._dwg.go import GoState
+
             return GoState(  # type:ignore
-                size=_states.size[_i],
-                chain_id_board=_states.chain_id_board[_i],
-                turn=_states.turn[_i],
+                size=_states.size[_i],  # type:ignore
+                chain_id_board=_states.chain_id_board[_i],  # type:ignore
+                turn=_states.turn[_i],  # type:ignore
             )
-        elif isinstance(_states, HexState):
+        elif _states.env_id == "hex":
+            from ._dwg.hex import HexState
+
             return HexState(
-                size=_states.size[_i],
-                turn=_states.turn[_i],
-                board=_states.board[_i],
+                size=_states.size[_i],  # type:ignore
+                turn=_states.turn[_i],  # type:ignore
+                board=_states.board[_i],  # type:ignore
             )
-        elif isinstance(_states, KuhnPokerState):
-            return KuhnPokerState(cards=_states.cards[_i], pot=_states.pot[_i])
-        elif isinstance(_states, LeducHoldemState):
+        elif _states.env_id == "kuhn_poker":
+            from ._dwg.kuhn_poker import KuhnPokerState
+
+            return KuhnPokerState(
+                cards=_states.cards[_i], pot=_states.pot[_i]  # type:ignore
+            )
+        elif _states.env_id == "leduc_holdem":
+            from ._dwg.leduc_holdem import LeducHoldemState
+
             return LeducHoldemState(
-                cards=_states.cards[_i],
-                chips=_states.chips[_i],
-                round=_states.round[_i],
+                cards=_states.cards[_i],  # type:ignore
+                chips=_states.chips[_i],  # type:ignore
+                round=_states.round[_i],  # type:ignore
             )
-        elif isinstance(_states, OthelloState):
+        elif _states.env_id == "othello":
+            from ._dwg.othello import OthelloState
+
             return OthelloState(
-                turn=_states.turn[_i],
-                board=_states.board[_i],
+                turn=_states.turn[_i],  # type:ignore
+                board=_states.board[_i],  # type:ignore
             )
-        elif isinstance(_states, Play2048State):
+        elif _states.env_id == "2048":
+            from ._dwg.play2048 import Play2048State
+
             return Play2048State(
-                board=_states.board[_i],
+                board=_states.board[_i],  # type:ignore
             )
-        elif isinstance(_states, ShogiState):
+        elif _states.env_id == "shogi":
+            from ._dwg.shogi import ShogiState
+
             return ShogiState(  # type:ignore
-                turn=_states.turn[_i],
-                piece_board=_states.piece_board[_i],
-                hand=_states.hand[_i],
+                turn=_states.turn[_i],  # type:ignore
+                piece_board=_states.piece_board[_i],  # type:ignore
+                hand=_states.hand[_i],  # type:ignore
             )
-        elif isinstance(_states, SparrowMahjongState):
+        elif _states.env_id == "sparrow_mahjong":
+            from ._dwg.sparrow_mahjong import SparrowMahjongState
+
             return SparrowMahjongState(
-                current_player=_states.current_player[_i],
-                turn=_states.turn[_i],
-                rivers=_states.rivers[_i],
-                hands=_states.hands[_i],
-                n_red_in_hands=_states.n_red_in_hands[_i],
-                is_red_in_river=_states.is_red_in_river[_i],
-                wall=_states.wall[_i],
-                draw_ix=_states.draw_ix[_i],
-                shuffled_players=_states.shuffled_players[_i],
-                dora=_states.dora[_i],
+                current_player=_states.current_player[_i],  # type:ignore
+                turn=_states.turn[_i],  # type:ignore
+                rivers=_states.rivers[_i],  # type:ignore
+                hands=_states.hands[_i],  # type:ignore
+                n_red_in_hands=_states.n_red_in_hands[_i],  # type:ignore
+                is_red_in_river=_states.is_red_in_river[_i],  # type:ignore
+                wall=_states.wall[_i],  # type:ignore
+                draw_ix=_states.draw_ix[_i],  # type:ignore
+                shuffled_players=_states.shuffled_players[_i],  # type:ignore
+                dora=_states.dora[_i],  # type:ignore
             )
-        elif isinstance(_states, TictactoeState):
+        elif _states.env_id == "tic_tac_toe":
+            from ._dwg.tictactoe import TictactoeState
+
             return TictactoeState(
-                current_player=_states.current_player[_i],
-                legal_action_mask=_states.legal_action_mask[_i],
-                terminated=_states.terminated[_i],
-                turn=_states.turn[_i],
-                board=_states.board[_i],
+                current_player=_states.current_player[_i],  # type:ignore
+                legal_action_mask=_states.legal_action_mask[_i],  # type:ignore
+                terminated=_states.terminated[_i],  # type:ignore
+                turn=_states.turn[_i],  # type:ignore
+                board=_states.board[_i],  # type:ignore
             )
         else:
             assert False
diff --git a/pgx/go.py b/pgx/go.py
--- a/pgx/go.py
+++ b/pgx/go.py
@@ -52,7 +52,10 @@ class State(core.State):
 
     @property
     def env_id(self) -> core.EnvId:
-        size = int(self.size.item())
+        try:
+            size = int(self.size.item())
+        except TypeError:
+            size = int(self.size[0].item())
         return f"go-{size}x{size}"  # type: ignore
 
 


Test patch: