Take screenshot when S is pressed while bot is playing

bot.py

@@ -1,13 +1,29 @@
 #!/usr/bin/env python3
+import os
 import time
+from itertools import count
 from strategy import State, moves_to_keys
 from interaction import get_exapunks_window, focus_window, \
-                        screenshot_board, press_keys
+                        screenshot_board, press_keys, listen_keys
+
+
+def save_screenshot(win):
+    board = screenshot_board(win).convert('RGB')
+    os.makedirs('screens', exist_ok=True)
+    for i in count(1):
+        path = 'screens/board%d.png' % i
+        if not os.path.exists(path):
+            print('save screenshot in', path)
+            board.save(path)
+            break
 
 
 if __name__ == '__main__':
     win = get_exapunks_window()
     focus_window(win)
+
+    listen_keys({'s': lambda: save_screenshot(win)})
+
     prev_score = None
 
     while True:

interaction.py

@@ -53,6 +53,7 @@ def screenshot_board(window):
 
 
 def press_keys(window, keys):
+    assert disp.has_extension('XTEST')
     for key in keys:
         keysym = XK.string_to_keysym(key)
         keycode = disp.keysym_to_keycode(keysym)
@@ -66,6 +67,17 @@ def press_keys(window, keys):
         time.sleep(KEY_DELAY)
 
 
+def listen_keys(handlers):
+    from pynput import keyboard
+
+    def handler(keycode):
+        key = str(keycode).replace('\'', '')
+        if key in handlers:
+            handlers[key]()
+
+    keyboard.Listener(on_release=handler).start()
+
+
 if __name__ == '__main__':
     win = get_exapunks_window()
     win.raise_window()