From e40fd90d5c76de84c52211e475e0a24d74955047 Mon Sep 17 00:00:00 2001
From: xs5871 <xs5871@qn-4.net>
Date: Fri, 17 Feb 2023 00:22:09 +0000
Subject: [PATCH] Make oneshot stay active across layer changes

---
 kmk/modules/oneshot.py | 10 ++++++---
 tests/test_oneshot.py  | 46 +++++++++++++++++++++++++++++++++++++-----
 2 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/kmk/modules/oneshot.py b/kmk/modules/oneshot.py
index 3f2e607..c86a095 100644
--- a/kmk/modules/oneshot.py
+++ b/kmk/modules/oneshot.py
@@ -1,5 +1,6 @@
 from kmk.keys import make_argumented_key
 from kmk.modules.holdtap import ActivationType, HoldTap, HoldTapKeyMeta
+from kmk.modules.layers import LayerKeyMeta
 from kmk.utils import Debug
 
 debug = Debug(__name__)
@@ -31,7 +32,9 @@ class OneShot(HoldTap):
             if key == current_key:
                 continue
 
-            if isinstance(current_key.meta, OneShotKeyMeta):
+            if (isinstance(current_key.meta, OneShotKeyMeta)) or (
+                isinstance(current_key.meta, LayerKeyMeta)
+            ):
                 keyboard.cancel_timeout(state.timeout_key)
                 if key.meta.tap_time is None:
                     tap_time = self.tap_time
@@ -50,13 +53,14 @@ class OneShot(HoldTap):
             elif state.activated == ActivationType.INTERRUPTED:
                 if is_pressed:
                     send_buffer = True
-                keyboard.set_timeout(0, lambda k=key: self.ht_released(k, keyboard))
+                self.key_buffer.insert(0, (0, key, False))
 
         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
 
+        self.send_key_buffer(keyboard)
+
         return current_key
 
     def osk_pressed(self, key, keyboard, *args, **kwargs):
diff --git a/tests/test_oneshot.py b/tests/test_oneshot.py
index 5774992..3cfc4ac 100644
--- a/tests/test_oneshot.py
+++ b/tests/test_oneshot.py
@@ -1,6 +1,7 @@
 import unittest
 
 from kmk.keys import KC
+from kmk.modules.layers import Layers
 from kmk.modules.oneshot import OneShot
 from tests.keyboard_test import KeyboardTest
 
@@ -8,17 +9,17 @@ from tests.keyboard_test import KeyboardTest
 class TestHoldTap(unittest.TestCase):
     def test_oneshot(self):
         keyboard = KeyboardTest(
-            [OneShot()],
+            [Layers(), OneShot()],
             [
                 [
-                    KC.A,
-                    KC.B,
+                    KC.OS(KC.MO(1), tap_time=50),
+                    KC.MO(1),
                     KC.C,
                     KC.D,
                     KC.OS(KC.E, tap_time=50),
                     KC.OS(KC.F, tap_time=50),
                 ],
-                [KC.N0, KC.N1, KC.N2, KC.N3, KC.N4],
+                [KC.N0, KC.N1, KC.N2, KC.N3, KC.OS(KC.LSFT, tap_time=50), KC.TRNS],
             ],
             debug_enabled=False,
         )
@@ -93,7 +94,7 @@ class TestHoldTap(unittest.TestCase):
                 (3, True),
                 (3, False),
             ],
-            [{KC.E}, {KC.E, KC.F}, {KC.E, KC.F, KC.D}, {KC.E, KC.F}, {KC.F}, {}],
+            [{KC.E}, {KC.E, KC.F}, {KC.E, KC.F, KC.D}, {KC.E, KC.F}, {KC.E}, {}],
         )
 
         keyboard.test(
@@ -109,3 +110,38 @@ class TestHoldTap(unittest.TestCase):
             ],
             [{KC.E}, {KC.E, KC.F}, {KC.E}, {}, {KC.D}, {}],
         )
+
+        keyboard.test(
+            'OS stacking with OS-layer',
+            [
+                (0, True),
+                (0, False),
+                (4, True),
+                (4, False),
+                (1, True),
+                (1, False),
+            ],
+            [{KC.LSFT}, {KC.LSFT, KC.N1}, {KC.LSFT}, {}],
+        )
+
+        keyboard.test(
+            'OS stacking with layer change',
+            [
+                (1, True),
+                (4, True),
+                (4, False),
+                (1, False),
+                (4, True),
+                (4, False),
+                (2, True),
+                (2, False),
+            ],
+            [
+                {KC.LSFT},
+                {KC.LSFT, KC.E},
+                {KC.LSFT, KC.E, KC.C},
+                {KC.LSFT, KC.E},
+                {KC.LSFT},
+                {},
+            ],
+        )