Fixup pynput import issues in headless tests

tests/test_cli.py

@@ -242,7 +242,7 @@ def test_annotate_data(annot_config, mocker):
     _ = mocker.patch("cv2.namedWindow")
     _ = mocker.patch("cv2.waitKey")
     _ = mocker.patch("cv2.imshow")
-    _ = mocker.patch("pynput.keyboard.Listener")
+    _ = mocker.patch("thelper.gui.utils.create_key_listener")
     fake_config = copy.deepcopy(annot_config)
     del fake_config["name"]
     with pytest.raises(AssertionError):

thelper/gui/__init__.py

@@ -12,5 +12,6 @@
 from thelper.gui.annotators import Annotator  # noqa: F401
 from thelper.gui.annotators import ImageSegmentAnnotator  # noqa: F401
 from thelper.gui.utils import create_annotator  # noqa: F401
+from thelper.gui.utils import create_key_listener  # noqa: F401
 
 logger = logging.getLogger("thelper.gui")

thelper/gui/annotators.py

@@ -301,8 +301,7 @@ def __init__(self, session_name, config, save_dir, datasets):
         self.image, self.mask = self.load(self.curr_sample_idx)
         self.image_display, self.image_display_base, self.mask_display = self.refresh_layers()
         self.gui_display = None
-        import pynput.keyboard
-        self.listener = pynput.keyboard.Listener(on_press=ImageSegmentAnnotator.on_press)
+        self.listener = thelper.gui.create_key_listener(ImageSegmentAnnotator.on_press)
         self.listener.start()
         cv.setMouseCallback(self.window_name, ImageSegmentAnnotator.on_mouse)
 

thelper/gui/utils.py

@@ -9,6 +9,12 @@
 logger = logging.getLogger(__name__)
 
 
+def create_key_listener(callback):
+    """Returns a key press listener based on pynput.keyboard (used for mocking)."""
+    import pynput.keyboard
+    return pynput.keyboard.Listener(on_press=callback)
+
+
 def create_annotator(session_name, save_dir, config, datasets):
     """Instantiates a GUI annotation tool based on the type contained in the config dictionary.