Implement use of joystick

Signed-off-by: Olivier Gayot <olivier.gayot@sigexec.com>

1 files changed, 41 insertions, 2 deletions

diff --git a/cameltris.py b/cameltris.py
index 5d1d9ba..58de00e 100755
--- a/cameltris.py
+++ b/cameltris.py
@@ -1,12 +1,22 @@
 #!/usr/bin/env python3
 
 import contextlib
+import enum
 import random
 import sys
 import time
 
 import pygame
 
+class PS3Controller(enum.Enum):
+    CROSS = 0
+    CIRCLE = 1
+    TRIANGLE = 2
+    SQUARE = 3
+    DOWN = 14
+    LEFT = 15
+    RIGHT = 16
+
 
 class WouldCollide(Exception):
     pass
@@ -354,6 +364,11 @@ das = 0
 
 pressing_down_countdown = None
 
+joysticks = [pygame.joystick.Joystick(x) for x in range(pygame.joystick.get_count())]
+for joystick in joysticks:
+    joystick.init()
+joystick = joysticks[0]
+
 while True:
     piece_drop_frames += 1
 
@@ -385,13 +400,37 @@ while True:
             if event.key in (pygame.K_DOWN, pygame.K_j):
                 pressing_down_countdown = None
 
+        elif event.type == pygame.JOYBUTTONDOWN:
+            with contextlib.suppress(WouldCollide):
+                if event.button == PS3Controller.RIGHT.value:
+                    move_piece_right()
+                    das = 0
+                if event.button == PS3Controller.LEFT.value:
+                    move_piece_left()
+                    das = 0
+                if event.button == PS3Controller.CROSS.value:
+                    rotate_piece_clockwise()
+                if event.button == PS3Controller.CIRCLE.value:
+                    rotate_piece_counter_clockwise()
+                if event.button == PS3Controller.DOWN.value:
+                    piece_drop_frames = 0
+                    pressing_down_countdown = 3
+                    try:
+                        move_piece_down()
+                    except WouldCollide:
+                        stick_piece()
+        elif event.type == pygame.JOYBUTTONUP:
+            if event.button == PS3Controller.DOWN.value:
+                pressing_down_countdown = None
+
     das += 1
     if das == 16:
         pressed_keys = pygame.key.get_pressed()
+        pressed_buttons = [joystick.get_button(x) for x in range(joystick.get_numbuttons())]
         with contextlib.suppress(WouldCollide):
-            if pressed_keys[pygame.K_RIGHT] or pressed_keys[pygame.K_l]:
+            if pressed_keys[pygame.K_RIGHT] or pressed_keys[pygame.K_l] or pressed_buttons[PS3Controller.RIGHT.value]:
                 move_piece_right()
-            if pressed_keys[pygame.K_LEFT] or pressed_keys[pygame.K_h]:
+            if pressed_keys[pygame.K_LEFT] or pressed_keys[pygame.K_h] or pressed_buttons[PS3Controller.LEFT.value]:
                 move_piece_left()
         das = 10