From 6532497bb2e295c92f6debb03ca2d721fa583eae Mon Sep 17 00:00:00 2001
From: xs5871 <xs5871@qn-4.net>
Date: Thu, 16 Feb 2023 05:59:46 +0000
Subject: [PATCH] Implement oneshot key stacking

---
 kmk/modules/oneshot.py | 44 ++++++++++++++++++++++++++++++------------
 tests/test_oneshot.py  | 30 +++++++++++++++++++++++++++-
 2 files changed, 61 insertions(+), 13 deletions(-)

diff --git a/kmk/modules/oneshot.py b/kmk/modules/oneshot.py
index 107fa1a..3f2e607 100644
--- a/kmk/modules/oneshot.py
+++ b/kmk/modules/oneshot.py
@@ -1,9 +1,13 @@
 from kmk.keys import make_argumented_key
 from kmk.modules.holdtap import ActivationType, HoldTap, HoldTapKeyMeta
+from kmk.utils import Debug
+
+debug = Debug(__name__)
 
 
-def oneshot_validator(kc, tap_time=None):
-    return HoldTapKeyMeta(tap=kc, hold=kc, prefer_hold=False, tap_time=tap_time)
+class OneShotKeyMeta(HoldTapKeyMeta):
+    def __init__(self, kc, tap_time=None):
+        super().__init__(tap=kc, hold=kc, prefer_hold=False, tap_time=tap_time)
 
 
 class OneShot(HoldTap):
@@ -12,30 +16,46 @@ class OneShot(HoldTap):
     def __init__(self):
         super().__init__()
         make_argumented_key(
-            validator=oneshot_validator,
+            validator=OneShotKeyMeta,
             names=('OS', 'ONESHOT'),
             on_press=self.osk_pressed,
             on_release=self.osk_released,
         )
 
     def process_key(self, keyboard, current_key, is_pressed, int_coord):
-        '''Release os key after interrupting keyup.'''
+        '''Release os key after interrupting non-os keyup, or reset timeout and
+        stack multiple os keys.'''
+        send_buffer = False
+
         for key, state in self.key_states.items():
             if key == current_key:
                 continue
 
+            if isinstance(current_key.meta, OneShotKeyMeta):
+                keyboard.cancel_timeout(state.timeout_key)
+                if key.meta.tap_time is None:
+                    tap_time = self.tap_time
+                else:
+                    tap_time = key.meta.tap_time
+                state.timeout_key = keyboard.set_timeout(
+                    tap_time,
+                    lambda k=key: self.on_tap_time_expired(k, keyboard),
+                )
+                continue
+
             if state.activated == ActivationType.PRESSED and is_pressed:
                 state.activated = ActivationType.HOLD_TIMEOUT
             elif state.activated == ActivationType.RELEASED and is_pressed:
                 state.activated = ActivationType.INTERRUPTED
             elif state.activated == ActivationType.INTERRUPTED:
                 if is_pressed:
-                    keyboard.remove_key(key.meta.tap)
-                    self.key_buffer.append((int_coord, current_key, is_pressed))
-                    keyboard.set_timeout(False, lambda: self.send_key_buffer(keyboard))
-                    current_key = None
-                else:
-                    self.ht_released(key, keyboard)
+                    send_buffer = True
+                keyboard.set_timeout(0, lambda k=key: self.ht_released(k, keyboard))
+
+        if send_buffer:
+            self.key_buffer.append((int_coord, current_key, is_pressed))
+            keyboard.set_timeout(0, lambda: self.send_key_buffer(keyboard))
+            current_key = None
 
         return current_key
 
@@ -51,8 +71,8 @@ class OneShot(HoldTap):
         try:
             state = self.key_states[key]
         except KeyError:
-            if keyboard.debug_enabled:
-                print(f'OneShot.osk_released: no such key {key}')
+            if debug.enabled:
+                debug(f'OneShot.osk_released: no such key {key}')
             return keyboard
 
         if state.activated == ActivationType.PRESSED:
diff --git a/tests/test_oneshot.py b/tests/test_oneshot.py
index 24558b5..5774992 100644
--- a/tests/test_oneshot.py
+++ b/tests/test_oneshot.py
@@ -67,7 +67,6 @@ class TestHoldTap(unittest.TestCase):
             [{KC.E}, {KC.D, KC.E}, {KC.E}, {}],
         )
 
-
         keyboard.test(
             'OS hold with multiple interrupt keys',
             [
@@ -81,3 +80,32 @@ class TestHoldTap(unittest.TestCase):
             ],
             [{KC.E}, {KC.D, KC.E}, {KC.E}, {KC.C, KC.E}, {KC.E}, {}],
         )
+
+        keyboard.test(
+            'OS stacking within timeout reset',
+            [
+                (4, True),
+                (4, False),
+                t_within,
+                (5, True),
+                (5, False),
+                t_within,
+                (3, True),
+                (3, False),
+            ],
+            [{KC.E}, {KC.E, KC.F}, {KC.E, KC.F, KC.D}, {KC.E, KC.F}, {KC.F}, {}],
+        )
+
+        keyboard.test(
+            'OS stacking timed out',
+            [
+                (4, True),
+                (4, False),
+                (5, True),
+                (5, False),
+                t_after,
+                (3, True),
+                (3, False),
+            ],
+            [{KC.E}, {KC.E, KC.F}, {KC.E}, {}, {KC.D}, {}],
+        )