Linux: better handling of the Xrandr extension (fixes BoboTiG#168)

BoboTiG · BoboTiG · commit 89a1a395f719 · 2020-06-30T17:36:31.000+02:00
diff --git a/CHANGELOG b/CHANGELOG
@@ -7,6 +7,7 @@ History:
       - fixed flake8 usage in pre-commit
       - the module is now available on conda (closes #170)
       - MSS: the implementation is now thread-safe on all OSes (fixes #169)
+      - Linux: better handling of the Xrandr extension (fixes #168)
       - tests: fixed a random bug on test_grab_with_tuple_percents() (fixes #142)
       - :heart: contributors: @
 
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -16,6 +16,7 @@ darwin.py
 
 linux.py
 --------
+ - Added ``MSS.has_extension()``
  - Renamed ``MSS.grab()`` to ``MSS._grab_impl()``
  - Renamed ``MSS.monitors`` to ``MSS._monitors_impl()``
 
diff --git a/mss/base.py b/mss/base.py
@@ -13,6 +13,7 @@
 from .tools import to_png
 
 if TYPE_CHECKING:
+    # pylint: disable=ungrouped-imports
     from typing import Any, Callable, Iterator, List, Optional, Type  # noqa
 
     from .models import Monitor, Monitors  # noqa
diff --git a/mss/linux.py b/mss/linux.py
@@ -10,7 +10,7 @@
 from types import SimpleNamespace
 from typing import TYPE_CHECKING
 
-from .base import MSSBase
+from .base import MSSBase, lock
 from .exception import ScreenShotError
 
 if TYPE_CHECKING:
@@ -227,10 +227,36 @@ def __init__(self, display=None):
 
         self.root = self.xlib.XDefaultRootWindow(self._get_display(display))
 
+        if not self.has_extension("RANDR"):
+            raise ScreenShotError("No Xrandr extension found.")
+
         # Fix for XRRGetScreenResources and XGetImage:
         #     expected LP_Display instance instead of LP_XWindowAttributes
         self.drawable = ctypes.cast(self.root, ctypes.POINTER(Display))
 
+    def has_extension(self, extension):
+        # type: (str) -> bool
+        """Return True if the given *extension* is part of the extensions list of the server."""
+        with lock:
+            byref = ctypes.byref
+            c_int = ctypes.c_int
+            major_opcode_return = c_int()
+            first_event_return = c_int()
+            first_error_return = c_int()
+
+            try:
+                self.xlib.XQueryExtension(
+                    self._get_display(),
+                    extension.encode("latin1"),
+                    byref(major_opcode_return),
+                    byref(first_event_return),
+                    byref(first_error_return),
+                )
+            except ScreenShotError:
+                return False
+            else:
+                return True
+
     def _get_display(self, disp=None):
         """
         Retrieve a thread-safe display from XOpenDisplay().
@@ -297,6 +323,11 @@ def cfactory(func, argtypes, restype, attr=self.xlib):
             pointer(XImage),
         )
         cfactory("XDestroyImage", [pointer(XImage)], void)
+        cfactory(
+            "XQueryExtension",
+            [pointer(Display), char_p, pointer(c_int), pointer(c_int), pointer(c_int)],
+            uint,
+        )
 
         # A simple benchmark calling 10 times those 2 functions:
         # XRRGetScreenResources():        0.1755971429956844 s
diff --git a/mss/tests/test_gnu_linux.py b/mss/tests/test_gnu_linux.py
@@ -121,3 +121,10 @@ def test_region_out_of_monitor_bounds():
         details = sct.get_error_details()
         assert details["xerror"]
         assert isinstance(details["xerror_details"], dict)
+
+
+def test_has_extension():
+    display = os.getenv("DISPLAY")
+    with mss.mss(display=display) as sct:
+        assert sct.has_extension("RANDR")
+        assert not sct.has_extension("NOEXT")
