fix(Linux): reset X server error handler on exit to prevent issues with Tk/Tkinter

Author: BoboTiG

CHANGELOG

@@ -5,6 +5,7 @@ History:
 8.0.0   2023/0x/xx
       - the whole source code was migrated to PEP 570 (Python Positional-Only Parameters)
       - removed support for Python 3.7
+      - Linux: reset X server error handler on exit to prevent issues with Tk/Tkinter (#235)
       - Linux: added mouse support (#232)
       - Linux: refactored how internal handles are stored to fix issues with multiple X servers, and TKinter.
                No more side effects, and when leaving the context manager, resources are all freed (#224, #234)

mss/linux.py

@@ -182,14 +182,32 @@ class XWindowAttributes(Structure):
 
 
 _ERROR = {}
+_X11 = find_library("X11")
+_XFIXES = find_library("Xfixes")
+_XRANDR = find_library("Xrandr")
+
+
+@CFUNCTYPE(c_int, POINTER(Display), POINTER(Event))
+def _default_error_handler(display: Display, event: Event) -> int:
+    """
+    Specifies the default program's supplied error handler.
+    It's useful when exiting MSS to prevent letting `_error_handler()` as default handler.
+    Doing so would crash when using Tk/Tkinter, see issue #220.
+
+    Interesting technical stuff can be found here:
+        https://core.tcl-lang.org/tk/file?name=generic/tkError.c&ci=a527ef995862cb50
+        https://github.com/tcltk/tk/blob/b9cdafd83fe77499ff47fa373ce037aff3ae286a/generic/tkError.c
+    """
+    # pylint: disable=unused-argument
+    return 0  # pragma: nocover
 
 
 @CFUNCTYPE(c_int, POINTER(Display), POINTER(Event))
 def _error_handler(display: Display, event: Event) -> int:
     """Specifies the program's supplied error handler."""
 
     # Get the specific error message
-    xlib = cdll.LoadLibrary(find_library("X11"))  # type: ignore[arg-type]
+    xlib = cdll.LoadLibrary(_X11)  # type: ignore[arg-type]
     get_error = xlib.XGetErrorText
     get_error.argtypes = [POINTER(Display), c_int, c_char_p, c_int]
     get_error.restype = c_void_p
@@ -278,24 +296,21 @@ def __init__(self, /, **kwargs: Any) -> None:
         if b":" not in display:
             raise ScreenShotError(f"Bad display value: {display!r}.")
 
-        x11 = find_library("X11")
-        if not x11:
+        if not _X11:
             raise ScreenShotError("No X11 library found.")
-        self.xlib = cdll.LoadLibrary(x11)
+        self.xlib = cdll.LoadLibrary(_X11)
 
         # Install the error handler to prevent interpreter crashes:
         # any error will raise a ScreenShotError exception.
         self.xlib.XSetErrorHandler(_error_handler)
 
-        xrandr = find_library("Xrandr")
-        if not xrandr:
+        if not _XRANDR:
             raise ScreenShotError("No Xrandr extension found.")
-        self.xrandr = cdll.LoadLibrary(xrandr)
+        self.xrandr = cdll.LoadLibrary(_XRANDR)
 
         if self.with_cursor:
-            xfixes = find_library("Xfixes")
-            if xfixes:
-                self.xfixes = cdll.LoadLibrary(xfixes)
+            if _XFIXES:
+                self.xfixes = cdll.LoadLibrary(_XFIXES)
             else:
                 self.with_cursor = False
 
@@ -314,10 +329,15 @@ def __init__(self, /, **kwargs: Any) -> None:
         self._handles.drawable = cast(self._handles.root, POINTER(Display))
 
     def close(self) -> None:
+        # Remove our error handler
+        self.xlib.XSetErrorHandler(_default_error_handler)
+
+        # Clean-up
         if self._handles.display is not None:
             self.xlib.XCloseDisplay(self._handles.display)
             self._handles.display = None
 
+        # Also empty the error dict
         _ERROR.clear()
 
     def _is_extension_enabled(self, name: str, /) -> bool:
@@ -350,7 +370,8 @@ def _set_cfunctions(self) -> None:
         }
         for func, (attr, argtypes, restype) in CFUNCTIONS.items():
             with suppress(AttributeError):
-                cfactory(attrs[attr], func, argtypes, restype, errcheck=_validate)
+                errcheck = None if func == "XSetErrorHandler" else _validate
+                cfactory(attrs[attr], func, argtypes, restype, errcheck=errcheck)
 
     def _monitors_impl(self) -> None:
         """Get positions of monitors. It will populate self._monitors."""

mss/tests/test_gnu_linux.py

@@ -2,7 +2,6 @@
 This is part of the MSS Python's module.
 Source: https://github.com/BoboTiG/python-mss
 """
-import ctypes.util
 import platform
 from unittest.mock import Mock, patch
 
@@ -88,30 +87,18 @@ def test_bad_display_structure(monkeypatch):
             pass
 
 
+@patch("mss.linux._X11", new=None)
 def test_no_xlib_library():
-    with patch("mss.linux.find_library", return_value=None):
-        with pytest.raises(ScreenShotError):
-            with mss.mss():
-                pass
+    with pytest.raises(ScreenShotError):
+        with mss.mss():
+            pass
 
 
+@patch("mss.linux._XRANDR", new=None)
 def test_no_xrandr_extension():
-    x11 = ctypes.util.find_library("X11")
-
-    def find_lib_mocked(lib):
-        """
-        Returns None to emulate no XRANDR library.
-        Returns the previous found X11 library else.
-
-        It is a naive approach, but works for now.
-        """
-
-        return None if lib == "Xrandr" else x11
-
-    # No `Xrandr` library
-    with patch("mss.linux.find_library", find_lib_mocked):
-        with pytest.raises(ScreenShotError):
-            mss.mss()
+    with pytest.raises(ScreenShotError):
+        with mss.mss():
+            pass
 
 
 @patch("mss.linux.MSS._is_extension_enabled", new=Mock(return_value=False))
@@ -190,23 +177,11 @@ def test_with_cursor(display: str):
     assert set(screenshot_with_cursor.rgb) == {0, 255}
 
 
+@patch("mss.linux._XFIXES", new=None)
 def test_with_cursor_but_not_xfixes_extension_found(display: str):
-    x11 = ctypes.util.find_library("X11")
-
-    def find_lib_mocked(lib):
-        """
-        Returns None to emulate no XRANDR library.
-        Returns the previous found X11 library else.
-
-        It is a naive approach, but works for now.
-        """
-
-        return None if lib == "Xfixes" else x11
-
-    with patch("mss.linux.find_library", find_lib_mocked):
-        with mss.mss(display=display, with_cursor=True) as sct:
-            assert not hasattr(sct, "xfixes")
-            assert not sct.with_cursor
+    with mss.mss(display=display, with_cursor=True) as sct:
+        assert not hasattr(sct, "xfixes")
+        assert not sct.with_cursor
 
 
 def test_with_cursor_failure(display: str):

mss/tests/test_issue_220.py

@@ -0,0 +1,47 @@
+"""
+This is part of the MSS Python's module.
+Source: https://github.com/BoboTiG/python-mss
+"""
+import pytest
+
+import mss
+
+tkinter = pytest.importorskip("tkinter")
+root = tkinter.Tk()
+
+
+def take_screenshot():
+    region = {"top": 370, "left": 1090, "width": 80, "height": 390}
+    with mss.mss() as sct:
+        sct.grab(region)
+
+
+def create_top_level_win():
+    top_level_win = tkinter.Toplevel(root)
+
+    take_screenshot_btn = tkinter.Button(top_level_win, text="Take screenshot", command=take_screenshot)
+    take_screenshot_btn.pack()
+
+    take_screenshot_btn.invoke()
+    root.update_idletasks()
+    root.update()
+
+    top_level_win.destroy()
+    root.update_idletasks()
+    root.update()
+
+
+def test_regression(capsys):
+    btn = tkinter.Button(root, text="Open TopLevel", command=create_top_level_win)
+    btn.pack()
+
+    # First screenshot: it works
+    btn.invoke()
+
+    # Second screenshot: it should work too
+    btn.invoke()
+
+    # Check there were no exceptions
+    captured = capsys.readouterr()
+    assert not captured.out
+    assert not captured.err