From d7eb293af873b55702879ec53e8225db94d8fd1d Mon Sep 17 00:00:00 2001
From: xs5871 <xs5871@qn-4.net>
Date: Sat, 8 Apr 2023 10:01:17 +0000
Subject: [PATCH] Implement autoshift module

---
 kmk/hid.py               |  9 ++++
 kmk/modules/autoshift.py | 83 ++++++++++++++++++++++++++++++++++++
 tests/test_autoshift.py  | 91 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 183 insertions(+)
 create mode 100644 kmk/modules/autoshift.py
 create mode 100644 tests/test_autoshift.py

diff --git a/kmk/hid.py b/kmk/hid.py
index 77c47d9..323b3d5 100644
--- a/kmk/hid.py
+++ b/kmk/hid.py
@@ -229,6 +229,15 @@ class AbstractHID:
         for idx in range(2, len(self._pd_report)):
             self._pd_report[idx] = 0x00
 
+    def has_key(self, key):
+        if isinstance(key, ModifierKey):
+            return bool(self.report_mods[0] & key.code)
+        else:
+            for code in self.report_non_mods:
+                if code == key.code:
+                    return True
+        return False
+
 
 class USBHID(AbstractHID):
     REPORT_BYTES = 9
diff --git a/kmk/modules/autoshift.py b/kmk/modules/autoshift.py
new file mode 100644
index 0000000..6a53e78
--- /dev/null
+++ b/kmk/modules/autoshift.py
@@ -0,0 +1,83 @@
+from kmk.keys import KC, Key
+from kmk.modules import Module
+from kmk.scheduler import cancel_task, create_task
+from kmk.utils import Debug
+
+debug = Debug(__name__)
+
+
+class Autoshift(Module):
+    def __init__(self, tap_time=300):
+        self.tap_time = tap_time
+
+        self._active = False
+        self._task = None
+        self._key = None
+
+    def during_bootup(self, keyboard):
+        self._task = create_task(lambda: self._shift(keyboard), after_ms=-1)
+
+    def before_matrix_scan(self, keyboard):
+        pass
+
+    def after_matrix_scan(self, keyboard):
+        pass
+
+    def process_key(self, keyboard, key, is_pressed, int_coord):
+        # Unshift on any key event
+        if self._active:
+            self._unshift(keyboard)
+            return key
+
+        # Only shift from an unshifted state
+        if keyboard._hid_helper.has_key(KC.LSHIFT):
+            return key
+
+        # Ignore rolls from tapped to hold
+        if not is_pressed and key is not self._key:
+            return key
+
+        # Only shift alpha keys, iff there's no pending potential shift
+        if (
+            is_pressed
+            and not self._key
+            and isinstance(key, Key)
+            and key.code
+            and KC.A.code <= key.code <= KC.Z.code
+        ):
+            create_task(self._task, after_ms=self.tap_time)
+            self._key = key
+        else:
+            cancel_task(self._task)
+            keyboard.resume_process_key(self, self._key, True)
+            if key is self._key:
+                keyboard.resume_process_key(self, self._key, False)
+            else:
+                keyboard.resume_process_key(self, key, True)
+            self._key = None
+
+    def before_hid_send(self, keyboard):
+        pass
+
+    def after_hid_send(self, keyboard):
+        pass
+
+    def on_powersave_enable(self, keyboard):
+        pass
+
+    def on_powersave_disable(self, keyboard):
+        pass
+
+    def _shift(self, keyboard):
+        if debug.enabled:
+            debug('activate')
+        self._active = True
+        keyboard.keys_pressed.add(KC.LSFT)
+        keyboard.resume_process_key(self, self._key, True)
+
+    def _unshift(self, keyboard):
+        if debug.enabled:
+            debug('deactivate')
+        self._active = False
+        self._key = None
+        keyboard.keys_pressed.remove(KC.LSFT)
diff --git a/tests/test_autoshift.py b/tests/test_autoshift.py
new file mode 100644
index 0000000..68e5dc8
--- /dev/null
+++ b/tests/test_autoshift.py
@@ -0,0 +1,91 @@
+import unittest
+
+from kmk.keys import KC
+from kmk.modules.autoshift import Autoshift
+from tests.keyboard_test import KeyboardTest
+
+tap_time = 3 * KeyboardTest.loop_delay_ms
+t_after = 4 * KeyboardTest.loop_delay_ms
+
+
+class TestAutoshift(unittest.TestCase):
+    def setUp(self):
+        self.kb = KeyboardTest(
+            [Autoshift(tap_time=tap_time)],
+            [
+                [
+                    KC.A,
+                    KC.N1,
+                    KC.HASH,
+                    KC.NO,
+                ]
+            ],
+            debug_enabled=False,
+        )
+
+    def test_tap_alpha(self):
+        self.kb.test(
+            '',
+            [(0, True), (0, False)],
+            [{KC.A}, {}],
+        )
+
+    def test_hold_alpha(self):
+        self.kb.test(
+            '',
+            [(0, True), t_after, (0, False)],
+            [{KC.A, KC.LSHIFT}, {}],
+        )
+
+    def test_hold_num(self):
+        self.kb.test(
+            '',
+            [(1, True), t_after, (1, False)],
+            [{KC.N1}, {}],
+        )
+
+    def test_hold_alpha_tap_num_within(self):
+        self.kb.test(
+            '',
+            [(0, True), (1, True), t_after, (1, False), (0, False)],
+            [{KC.A}, {KC.A, KC.N1}, {KC.A}, {}],
+        )
+
+    def test_hold_alpha_tap_num_after(self):
+        self.kb.test(
+            '',
+            [(0, True), t_after, (1, True), (1, False), (0, False)],
+            [{KC.A, KC.LSHIFT}, {KC.A, KC.N1}, {KC.A}, {}],
+        )
+
+    def test_hold_num_hold_alpha(self):
+        self.kb.test(
+            '',
+            [(1, True), (0, True), t_after, (0, False), (1, False)],
+            [{KC.N1}, {KC.N1, KC.A, KC.LSHIFT}, {KC.N1}, {}],
+        )
+
+    def test_roll_num_hold_alpha(self):
+        self.kb.test(
+            '',
+            [(1, True), (0, True), (1, False), t_after, (0, False)],
+            [{KC.N1}, {}, {KC.A, KC.LSHIFT}, {}],
+        )
+
+    def test_hold_shifted_hold_alpha(self):
+        self.kb.test(
+            '',
+            [(2, True), (0, True), t_after, (2, False), (0, False)],
+            [{KC.LSHIFT, KC.HASH}, {KC.LSHIFT, KC.HASH, KC.A}, {KC.A}, {}],
+        )
+
+    def test_hold_internal(self):
+        self.kb.test(
+            '',
+            [(3, True), t_after, (3, False)],
+            [],
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()