remove duplicated code

connectfour/app.py

@@ -12,6 +12,7 @@ demo = gr.Blocks()
 
 with demo:
     connect4 = Connect4("You")
+    drop_token_btns = []
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -33,20 +34,9 @@ with demo:
             )
 
             with gr.Row():
-                with gr.Column(scale=1, min_width=20):
-                    drop_token0_btn = gr.Button("X")
-                with gr.Column(scale=1, min_width=20):
-                    drop_token1_btn = gr.Button("X")
-                with gr.Column(scale=1, min_width=20):
-                    drop_token2_btn = gr.Button("X")
-                with gr.Column(scale=1, min_width=20):
-                    drop_token3_btn = gr.Button("X")
-                with gr.Column(scale=1, min_width=20):
-                    drop_token4_btn = gr.Button("X")
-                with gr.Column(scale=1, min_width=20):
-                    drop_token5_btn = gr.Button("X")
-                with gr.Column(scale=1, min_width=20):
-                    drop_token6_btn = gr.Button("X")
+                for i in range(7):
+                    with gr.Column(min_width=20):
+                        drop_token_btns.append(gr.Button("X", elem_id=i))
 
     who_plays_first.change(
         connect4.init_env, who_plays_first, outputs=[output, game_state]
@@ -62,7 +52,7 @@ with demo:
 
     def wait(game_state_value):
         if game_state_value == "Game On":
-            time.sleep(1)
+            time.sleep(0.7)
             return gr.Checkbox.update(interactive=False)
         else:
             return gr.Checkbox.update(interactive=True)
@@ -79,89 +69,28 @@ with demo:
             gr.Checkbox.update(interactive=True),
         )
 
-    drop_token0_btn.click(
-        lambda: connect4.play(0),
-        outputs=[output, game_state],
-    ).then(
-        wait, inputs=[game_state], outputs=who_plays_first
-    ).then(bot, inputs=[game_state], outputs=[output, game_state, who_plays_first])
-
-    drop_token1_btn.click(
-        lambda: connect4.play(1),
-        outputs=[output, game_state],
-    ).then(
-        wait, inputs=[game_state], outputs=who_plays_first
-    ).then(bot, inputs=[game_state], outputs=[output, game_state, who_plays_first])
-
-    drop_token2_btn.click(
-        lambda: connect4.play(2),
-        outputs=[output, game_state],
-    ).then(
-        wait, inputs=[game_state], outputs=who_plays_first
-    ).then(bot, inputs=[game_state], outputs=[output, game_state, who_plays_first])
-
-    drop_token3_btn.click(
-        lambda: connect4.play(3),
-        outputs=[output, game_state],
-    ).then(
-        wait, inputs=[game_state], outputs=who_plays_first
-    ).then(bot, inputs=[game_state], outputs=[output, game_state, who_plays_first])
-
-    drop_token4_btn.click(
-        lambda: connect4.play(4),
-        outputs=[output, game_state],
-    ).then(
-        wait, inputs=[game_state], outputs=who_plays_first
-    ).then(bot, inputs=[game_state], outputs=[output, game_state, who_plays_first])
-
-    drop_token5_btn.click(
-        lambda: connect4.play(5),
-        outputs=[output, game_state],
-    ).then(
-        wait, inputs=[game_state], outputs=who_plays_first
-    ).then(bot, inputs=[game_state], outputs=[output, game_state, who_plays_first])
-
-    drop_token6_btn.click(
-        lambda: connect4.play(6),
-        outputs=[output, game_state],
-    ).then(
-        wait, inputs=[game_state], outputs=who_plays_first
-    ).then(bot, inputs=[game_state], outputs=[output, game_state, who_plays_first])
+    def click_column(evt: gr.EventData):
+        output, game_state = connect4.play(int(evt.target.elem_id))
+        return output, game_state
 
     def game_state_change(value):
         if value == "Game On":
-            return [
-                gr.Button.update(interactive=True),
-                gr.Button.update(interactive=True),
-                gr.Button.update(interactive=True),
-                gr.Button.update(interactive=True),
-                gr.Button.update(interactive=True),
-                gr.Button.update(interactive=True),
-                gr.Button.update(interactive=True),
-            ]
+            return [gr.Button.update(interactive=True)] * 7
         else:
-            return [
-                gr.Button.update(interactive=False),
-                gr.Button.update(interactive=False),
-                gr.Button.update(interactive=False),
-                gr.Button.update(interactive=False),
-                gr.Button.update(interactive=False),
-                gr.Button.update(interactive=False),
-                gr.Button.update(interactive=False),
-            ]
+            return [gr.Button.update(interactive=False)] * 7
+
+    for i in range(7):
+        drop_token_btns[i].click(
+            click_column,
+            outputs=[output, game_state],
+        ).then(
+            wait, inputs=[game_state], outputs=who_plays_first
+        ).then(bot, inputs=[game_state], outputs=[output, game_state, who_plays_first])
 
     game_state.change(
         game_state_change,
         game_state,
-        outputs=[
-            drop_token0_btn,
-            drop_token1_btn,
-            drop_token2_btn,
-            drop_token3_btn,
-            drop_token4_btn,
-            drop_token5_btn,
-            drop_token6_btn,
-        ],
+        outputs=drop_token_btns,
     )
 
 

connectfour/connect4.py

@@ -30,28 +30,28 @@ class Connect4:
         return flatten_obs[np.newaxis, ...].astype(np.float32)
 
     def play(self, action=None):
-        if self.has_erroneous_state():
+        try:
+            if self.human != self.current_player_id:
+                action = self.session.run(
+                    ["output"],
+                    {
+                        "obs": self.flatten_observation(
+                            self.env.observe(self.current_player_id)
+                        ),
+                        "state_ins": [],
+                    },
+                )
+                action = int(np.argmax(action[0]))
+
+            if action not in self.legal_moves:
+                action = np.random.choice(self.legal_moves)
+
+            self.env.step(action)
+
+            return self.render_and_state
+        except:
             return self.blue_screen()
 
-        if self.human != self.current_player_id:
-            action = self.session.run(
-                ["output"],
-                {
-                    "obs": self.flatten_observation(
-                        self.env.observe(self.current_player_id)
-                    ),
-                    "state_ins": [],
-                },
-            )
-            action = int(np.argmax(action[0]))
-
-        if action not in self.legal_moves:
-            action = np.random.choice(self.legal_moves)
-
-        self.env.step(action)
-
-        return self.render_and_state
-
     @property
     def current_player_id(self):
         return self.env.agent_selection
@@ -90,9 +90,3 @@ class Connect4:
             error_screen = np.load(f)
 
         return (error_screen, "Restart the Game")
-
-    def has_erroneous_state(self):
-        obs = self.env.observe(self.current_player_id)
-        if len(list(obs.keys())) == 0:
-            return True
-        return False