Implement an XShmGetImage-based backend.

This is close to complete, but there's a few things that need to be
chased down: notably, the test_thread_safety test is failing, for some
reason.  It also currently doesn't work correctly if the root window
size increases.

That said, this is quite promising: on my computer, the new backend
can take 4k screenshots at 30-34 fps, while the XGetImage backend
could only run at 11-14 fps.

Author: jholveck

src/mss/linux/__init__.py

@@ -14,6 +14,10 @@ def mss(backend: str = "default", **kwargs: Any) -> MSSBase:
         from . import xgetimage  # noqa: PLC0415
 
         return xgetimage.MSS(**kwargs)
+    if backend == "xshmgetimage":
+        from . import xshmgetimage  # noqa: PLC0415
+
+        return xshmgetimage.MSS(**kwargs)
     msg = f"Backend {backend!r} not (yet?) implemented."
     raise ScreenShotError(msg)
 

src/mss/linux/base.py

@@ -9,6 +9,7 @@
 from .xcb import LIB
 
 if TYPE_CHECKING:
+    from mss.models import Monitor
     from mss.screenshot import ScreenShot
 
 SUPPORTED_DEPTHS = {24, 32}
@@ -48,8 +49,8 @@ def __init__(self, /, **kwargs: Any) -> None:  # noqa: PLR0912
         # Get the connection setup information that was included when we connected.
         xcb_setup = xcb.get_setup(self.conn)
         screens = xcb.setup_roots(xcb_setup)
-        pref_screen = screens[pref_screen_num]
-        self.root = self.drawable = pref_screen.root
+        self.pref_screen = screens[pref_screen_num]
+        self.root = self.drawable = self.pref_screen.root
 
         # We don't probe the XFixes presence or version until we need it.
         self._xfixes_ready: bool | None = None
@@ -61,8 +62,8 @@ def __init__(self, /, **kwargs: Any) -> None:  # noqa: PLR0912
         # Currently, we assume that the drawable we're capturing is the root; when we add single-window capture,
         # we'll have to ask the server for its depth and visual.
         assert self.root == self.drawable  # noqa: S101
-        self.drawable_depth = pref_screen.root_depth
-        self.drawable_visual_id = pref_screen.root_visual.value
+        self.drawable_depth = self.pref_screen.root_depth
+        self.drawable_visual_id = self.pref_screen.root_visual.value
         # Server image byte order
         if xcb_setup.image_byte_order != xcb.ImageOrder.LSBFirst:
             msg = "Only X11 servers using LSB-First images are supported."
@@ -92,7 +93,7 @@ def __init__(self, /, **kwargs: Any) -> None:  # noqa: PLR0912
             raise ScreenShotError(msg)
         # Visual, the interpretation of pixels (like indexed, grayscale, etc).  (Visuals are arranged by depth, so
         # we iterate over the depths first.)
-        for xcb_depth in xcb.screen_allowed_depths(pref_screen):
+        for xcb_depth in xcb.screen_allowed_depths(self.pref_screen):
             if xcb_depth.depth == self.drawable_depth:
                 break
         else:
@@ -226,3 +227,41 @@ def _cursor_impl(self) -> ScreenShot:
         # unfortunate historical accident that makes it have to return the cursor image in a different format.
 
         return self.cls_image(data, region)
+
+    def _grab_impl_xgetimage(self, monitor: Monitor, /) -> ScreenShot:
+        """Retrieve all pixels from a monitor using GetImage.
+
+        This is used by the XGetImage backend, and also the XShmGetImage
+        backend in fallback mode.
+        """
+
+        if self.conn is None:
+            msg = "Cannot take screenshot while the connection is closed"
+            raise ScreenShotError(msg)
+
+        img_reply = xcb.get_image(
+            self.conn,
+            xcb.ImageFormat.ZPixmap,
+            self.drawable,
+            monitor["left"],
+            monitor["top"],
+            monitor["width"],
+            monitor["height"],
+            ALL_PLANES,
+        )
+
+        # Now, save the image.  This is a reference into the img_reply structure.
+        img_data_arr = xcb.get_image_data(img_reply)
+        # Copy this into a new bytearray, so that it will persist after we clear the image structure.
+        img_data = bytearray(img_data_arr)
+
+        if img_reply.depth != self.drawable_depth or img_reply.visual.value != self.drawable_visual_id:
+            # This should never happen; a window can't change its visual.
+            msg = (
+                "Server returned an image with a depth or visual different than it initially reported: "
+                f"expected {self.drawable_depth},{hex(self.drawable_visual_id)}, "
+                f"got {img_reply.depth},{hex(img_reply.visual.value)}"
+            )
+            raise ScreenShotError(msg)
+
+        return self.cls_image(img_data, monitor)

src/mss/linux/xcb.py

@@ -11,6 +11,8 @@
     RANDR_MINOR_VERSION,
     RENDER_MAJOR_VERSION,
     RENDER_MINOR_VERSION,
+    SHM_MAJOR_VERSION,
+    SHM_MINOR_VERSION,
     XFIXES_MAJOR_VERSION,
     XFIXES_MINOR_VERSION,
     Atom,
@@ -52,6 +54,10 @@
     ScreenIterator,
     Setup,
     SetupIterator,
+    ShmCreateSegmentReply,
+    ShmGetImageReply,
+    ShmQueryVersionReply,
+    ShmSeg,
     Timestamp,
     VisualClass,
     Visualid,
@@ -91,6 +97,12 @@
     setup_pixmap_formats,
     setup_roots,
     setup_vendor,
+    shm_attach_fd,
+    shm_create_segment,
+    shm_create_segment_reply_fds,
+    shm_detach,
+    shm_get_image,
+    shm_query_version,
     xfixes_get_cursor_image,
     xfixes_get_cursor_image_cursor_image,
     xfixes_query_version,

src/mss/linux/xcbgen.py

@@ -631,7 +631,7 @@ def xfixes_get_cursor_image_cursor_image(r: XfixesGetCursorImageReply) -> Array[
     )
 
 
-def shm_create_segment_reply_fds(c: Connection, r: ShmCreateSegmentReply) -> _Pointer[c_int]:
+def shm_create_segment_reply_fds(c: Connection | _Pointer[Connection], r: ShmCreateSegmentReply) -> _Pointer[c_int]:
     return LIB.shm.xcb_shm_create_segment_reply_fds(c, r)
 
 
@@ -665,7 +665,7 @@ def get_property(
 
 
 def no_operation(c: Connection) -> None:
-    return LIB.xcb.xcb_no_operation(c).check(c)
+    return LIB.xcb.xcb_no_operation_checked(c).check(c)
 
 
 def randr_query_version(
@@ -716,7 +716,7 @@ def shm_get_image(
 
 
 def shm_attach_fd(c: Connection, shmseg: ShmSeg, shm_fd: c_int | int, read_only: c_uint8 | int) -> None:
-    return LIB.shm.xcb_shm_attach_fd(c, shmseg, shm_fd, read_only).check(c)
+    return LIB.shm.xcb_shm_attach_fd_checked(c, shmseg, shm_fd, read_only).check(c)
 
 
 def shm_create_segment(
@@ -726,7 +726,7 @@ def shm_create_segment(
 
 
 def shm_detach(c: Connection, shmseg: ShmSeg) -> None:
-    return LIB.shm.xcb_shm_detach(c, shmseg).check(c)
+    return LIB.shm.xcb_shm_detach_checked(c, shmseg).check(c)
 
 
 def xfixes_query_version(
@@ -860,8 +860,8 @@ def initialize() -> None:  # noqa: PLR0915
         [POINTER(Connection), c_uint8, Window, Atom, Atom, c_uint32, c_uint32],
         GetPropertyReply,
     )
-    LIB.xcb.xcb_no_operation.argtypes = (Connection,)
-    LIB.xcb.xcb_no_operation.restype = VoidCookie
+    LIB.xcb.xcb_no_operation_checked.argtypes = (POINTER(Connection),)
+    LIB.xcb.xcb_no_operation_checked.restype = VoidCookie
     initialize_xcb_typed_func(
         LIB.randr, "xcb_randr_query_version", [POINTER(Connection), c_uint32, c_uint32], RandrQueryVersionReply
     )
@@ -890,21 +890,21 @@ def initialize() -> None:  # noqa: PLR0915
         [POINTER(Connection), Drawable, c_int16, c_int16, c_uint16, c_uint16, c_uint32, c_uint8, ShmSeg, c_uint32],
         ShmGetImageReply,
     )
-    LIB.shm.xcb_shm_attach_fd.argtypes = (
-        Connection,
+    LIB.shm.xcb_shm_attach_fd_checked.argtypes = (
+        POINTER(Connection),
         ShmSeg,
         c_int,
         c_uint8,
     )
-    LIB.shm.xcb_shm_attach_fd.restype = VoidCookie
+    LIB.shm.xcb_shm_attach_fd_checked.restype = VoidCookie
     initialize_xcb_typed_func(
         LIB.shm, "xcb_shm_create_segment", [POINTER(Connection), ShmSeg, c_uint32, c_uint8], ShmCreateSegmentReply
     )
-    LIB.shm.xcb_shm_detach.argtypes = (
-        Connection,
+    LIB.shm.xcb_shm_detach_checked.argtypes = (
+        POINTER(Connection),
         ShmSeg,
     )
-    LIB.shm.xcb_shm_detach.restype = VoidCookie
+    LIB.shm.xcb_shm_detach_checked.restype = VoidCookie
     initialize_xcb_typed_func(
         LIB.xfixes, "xcb_xfixes_query_version", [POINTER(Connection), c_uint32, c_uint32], XfixesQueryVersionReply
     )

src/mss/linux/xgetimage.py

@@ -1,9 +1,7 @@
-from mss.exception import ScreenShotError
 from mss.models import Monitor
 from mss.screenshot import ScreenShot
 
-from . import xcb
-from .base import ALL_PLANES, MSSXCBBase
+from .base import MSSXCBBase
 
 
 class MSS(MSSXCBBase):
@@ -15,36 +13,6 @@ class MSS(MSSXCBBase):
     * XFixes: Including the cursor.
     """
 
-    def _grab_impl(self, monitor: Monitor, /) -> ScreenShot:
+    def _grab_impl(self, monitor: Monitor) -> ScreenShot:
         """Retrieve all pixels from a monitor. Pixels have to be RGBX."""
-
-        if self.conn is None:
-            msg = "Cannot take screenshot while the connection is closed"
-            raise ScreenShotError(msg)
-
-        img_reply = xcb.get_image(
-            self.conn,
-            xcb.ImageFormat.ZPixmap,
-            self.drawable,
-            monitor["left"],
-            monitor["top"],
-            monitor["width"],
-            monitor["height"],
-            ALL_PLANES,
-        )
-
-        # Now, save the image.  This is a reference into the img_reply structure.
-        img_data_arr = xcb.get_image_data(img_reply)
-        # Copy this into a new bytearray, so that it will persist after we clear the image structure.
-        img_data = bytearray(img_data_arr)
-
-        if img_reply.depth != self.drawable_depth or img_reply.visual.value != self.drawable_visual_id:
-            # This should never happen; a window can't change its visual.
-            msg = (
-                "Server returned an image with a depth or visual different than it initially reported: "
-                f"expected {self.drawable_depth},{hex(self.drawable_visual_id)}, "
-                f"got {img_reply.depth},{hex(img_reply.visual.value)}"
-            )
-            raise ScreenShotError(msg)
-
-        return self.cls_image(img_data, monitor)
+        return super()._grab_impl_xgetimage(monitor)