Add XCB MIT-SHM support, and factor out the XCB setup

This only adds the support for the XCB MIT-SHM extension to mss's
internal xcb libraries. The actual usage of shared memory for
screenshots will be done in a future commit.

Author: jholveck

src/mss/linux/base.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from mss.base import MSSBase
+from mss.exception import ScreenShotError
+
+from . import xcb
+from .xcb import LIB
+
+if TYPE_CHECKING:
+    from mss.screenshot import ScreenShot
+
+SUPPORTED_DEPTHS = {24, 32}
+SUPPORTED_BITS_PER_PIXEL = 32
+SUPPORTED_RED_MASK = 0xFF0000
+SUPPORTED_GREEN_MASK = 0x00FF00
+SUPPORTED_BLUE_MASK = 0x0000FF
+ALL_PLANES = 0xFFFFFFFF  # XCB doesn't define AllPlanes
+
+
+class MSSXCBBase(MSSBase):
+    """Base class for XCB-based screenshot implementations.
+
+    This class provides common XCB initialization and monitor detection logic
+    that can be shared across different XCB screenshot methods (XGetImage,
+    XShmGetImage, XComposite, etc.).
+    """
+
+    def __init__(self, /, **kwargs: Any) -> None:  # noqa: PLR0912
+        """Initialize an XCB connection and validate the display configuration.
+
+        Args:
+            **kwargs: Keyword arguments, including optional 'display' for X11 display string.
+
+        Raises:
+            ScreenShotError: If the display configuration is not supported.
+        """
+        super().__init__(**kwargs)
+
+        display = kwargs.get("display", b"")
+        if not display:
+            display = None
+
+        self.conn: xcb.Connection | None
+        self.conn, pref_screen_num = xcb.connect(display)
+
+        # 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
+
+        # We don't probe the XFixes presence or version until we need it.
+        self._xfixes_ready: bool | None = None
+
+        # Probe the visuals (and related information), and make sure that our drawable is in an acceptable format.
+        # These iterations and tests don't involve any traffic with the server; it's all stuff that was included in
+        # the connection setup.  Effectively all modern setups will be acceptable, but we verify to be sure.
+
+        # 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
+        # Server image byte order
+        if xcb_setup.image_byte_order != xcb.ImageOrder.LSBFirst:
+            msg = "Only X11 servers using LSB-First images are supported."
+            raise ScreenShotError(msg)
+        # Depth
+        if self.drawable_depth not in SUPPORTED_DEPTHS:
+            msg = f"Only screens of color depth 24 or 32 are supported, not {self.drawable_depth}"
+            raise ScreenShotError(msg)
+        # Format (i.e., bpp, padding)
+        for format_ in xcb.setup_pixmap_formats(xcb_setup):
+            if format_.depth == self.drawable_depth:
+                break
+        else:
+            msg = f"Internal error: drawable's depth {self.drawable_depth} not found in screen's supported formats"
+            raise ScreenShotError(msg)
+        drawable_format = format_
+        if drawable_format.bits_per_pixel != SUPPORTED_BITS_PER_PIXEL:
+            msg = (
+                f"Only screens at 32 bpp (regardless of color depth) are supported; "
+                f"got {drawable_format.bits_per_pixel} bpp"
+            )
+            raise ScreenShotError(msg)
+        if drawable_format.scanline_pad != SUPPORTED_BITS_PER_PIXEL:
+            # To clarify the padding: the scanline_pad is the multiple that the scanline gets padded to.  If there
+            # is no padding, then it will be the same as one pixel's size.
+            msg = "Screens with scanline padding are not supported"
+            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):
+            if xcb_depth.depth == self.drawable_depth:
+                break
+        else:
+            msg = "Internal error: drawable's depth not found in screen's supported depths"
+            raise ScreenShotError(msg)
+        for visual_info in xcb.depth_visuals(xcb_depth):
+            if visual_info.visual_id.value == self.drawable_visual_id:
+                break
+        else:
+            msg = "Internal error: drawable's visual not found in screen's supported visuals"
+            raise ScreenShotError(msg)
+        if visual_info.class_ not in {xcb.VisualClass.TrueColor, xcb.VisualClass.DirectColor}:
+            msg = "Only TrueColor and DirectColor visuals are supported"
+            raise ScreenShotError(msg)
+        if (
+            visual_info.red_mask != SUPPORTED_RED_MASK
+            or visual_info.green_mask != SUPPORTED_GREEN_MASK
+            or visual_info.blue_mask != SUPPORTED_BLUE_MASK
+        ):
+            # There are two ways to phrase this layout: BGRx accounts for the byte order, while xRGB implies the
+            # native word order.  Since we return the data as a byte array, we use the former.  By the time we get
+            # to this point, we've already checked the endianness and depth, so this is pretty much never going to
+            # happen anyway.
+            msg = "Only visuals with BGRx ordering are supported"
+            raise ScreenShotError(msg)
+
+    def close(self) -> None:
+        """Close the XCB connection."""
+        if self.conn is not None:
+            xcb.disconnect(self.conn)
+        self.conn = None
+
+    def _monitors_impl(self) -> None:
+        """Get positions of monitors. It will populate self._monitors."""
+        if self.conn is None:
+            msg = "Cannot identify monitors while the connection is closed"
+            raise ScreenShotError(msg)
+
+        # The first entry is the whole X11 screen that the root is on.  That's the one that covers all the
+        # monitors.
+        root_geom = xcb.get_geometry(self.conn, self.root)
+        self._monitors.append(
+            {
+                "left": root_geom.x,
+                "top": root_geom.y,
+                "width": root_geom.width,
+                "height": root_geom.height,
+            }
+        )
+
+        # After that, we have one for each monitor on that X11 screen.  For decades, that's been handled by
+        # Xrandr. We don't presently try to work with Xinerama.  So, we're going to check the different outputs,
+        # according to Xrandr.  If that fails, we'll just leave the one root covering everything.
+
+        # Make sure we have the Xrandr extension we need.  This will query the cache that we started populating in
+        # __init__.
+        randr_ext_data = xcb.get_extension_data(self.conn, LIB.randr_id)
+        if not randr_ext_data.present:
+            return
+
+        # We ask the server to give us anything up to the version we support (i.e., what we expect the reply
+        # structs to look like).  If the server only supports 1.2, then that's what it'll give us, and we're ok
+        # with that, but we also use a faster path if the server implements at least 1.3.
+        randr_version_data = xcb.randr_query_version(self.conn, xcb.RANDR_MAJOR_VERSION, xcb.RANDR_MINOR_VERSION)
+        randr_version = (randr_version_data.major_version, randr_version_data.minor_version)
+        if randr_version < (1, 2):
+            return
+
+        screen_resources: xcb.RandrGetScreenResourcesReply | xcb.RandrGetScreenResourcesCurrentReply
+        # Check to see if we have the xcb_randr_get_screen_resources_current function in libxcb-randr, and that
+        # the server supports it.
+        if hasattr(LIB.randr, "xcb_randr_get_screen_resources_current") and randr_version >= (1, 3):
+            screen_resources = xcb.randr_get_screen_resources_current(self.conn, self.drawable.value)
+            crtcs = xcb.randr_get_screen_resources_current_crtcs(screen_resources)
+        else:
+            # Either the client or the server doesn't support the _current form.  That's ok; we'll use the old
+            # function, which forces a new query to the physical monitors.
+            screen_resources = xcb.randr_get_screen_resources(self.conn, self.drawable)
+            crtcs = xcb.randr_get_screen_resources_crtcs(screen_resources)
+
+        for crtc in crtcs:
+            crtc_info = xcb.randr_get_crtc_info(self.conn, crtc, screen_resources.config_timestamp)
+            if crtc_info.num_outputs == 0:
+                continue
+            self._monitors.append(
+                {"left": crtc_info.x, "top": crtc_info.y, "width": crtc_info.width, "height": crtc_info.height}
+            )
+
+        # Extra credit would be to enumerate the virtual desktops; see
+        # https://specifications.freedesktop.org/wm/latest/ar01s03.html.  But I don't know how widely-used that
+        # style is.
+
+    def _cursor_impl_check_xfixes(self) -> bool:
+        if self.conn is None:
+            msg = "Cannot take screenshot while the connection is closed"
+            raise ScreenShotError(msg)
+
+        xfixes_ext_data = xcb.get_extension_data(self.conn, LIB.xfixes_id)
+        if not xfixes_ext_data.present:
+            return False
+
+        reply = xcb.xfixes_query_version(self.conn, xcb.XFIXES_MAJOR_VERSION, xcb.XFIXES_MINOR_VERSION)
+        # We can work with 2.0 and later, but not sure about the actual minimum version we can use.  That's ok;
+        # everything these days is much more modern.
+        return (reply.major_version, reply.minor_version) >= (2, 0)
+
+    def _cursor_impl(self) -> ScreenShot:
+        """Retrieve all cursor data. Pixels have to be RGBx."""
+
+        if self.conn is None:
+            msg = "Cannot take screenshot while the connection is closed"
+            raise ScreenShotError(msg)
+
+        if self._xfixes_ready is None:
+            self._xfixes_ready = self._cursor_impl_check_xfixes()
+        if not self._xfixes_ready:
+            msg = "Server does not have XFixes, or the version is too old."
+            raise ScreenShotError(msg)
+
+        cursor_img = xcb.xfixes_get_cursor_image(self.conn)
+        region = {
+            "left": cursor_img.x - cursor_img.xhot,
+            "top": cursor_img.y - cursor_img.yhot,
+            "width": cursor_img.width,
+            "height": cursor_img.height,
+        }
+
+        data_arr = xcb.xfixes_get_cursor_image_cursor_image(cursor_img)
+        data = bytearray(data_arr)
+        # We don't need to do the same array slice-and-dice work as the Xlib-based implementation: Xlib has an
+        # unfortunate historical accident that makes it have to return the cursor image in a different format.
+
+        return self.cls_image(data, region)

src/mss/linux/xcb.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
-from ctypes import Structure, c_int, c_uint8, c_uint16, c_uint32
+from ctypes import Structure, _Pointer, c_int, c_uint8, c_uint16, c_uint32
 
 from . import xcbgen
 
 # We import these just so they're re-exported to our users.
-# ruff: noqa: F401
+# ruff: noqa: F401, TC001
 from .xcbgen import (
     RANDR_MAJOR_VERSION,
     RANDR_MINOR_VERSION,
@@ -97,7 +97,7 @@
 )
 
 # These are also here to re-export.
-from .xcbhelpers import LIB, Connection, XError
+from .xcbhelpers import LIB, XID, Connection, QueryExtensionReply, XcbExtension, XError
 
 XCB_CONN_ERROR = 1
 XCB_CONN_CLOSED_EXT_NOTSUPPORTED = 2
@@ -120,6 +120,53 @@
 }
 
 
+#### High-level XCB function wrappers
+
+
+def get_extension_data(
+    xcb_conn: Connection | _Pointer[Connection], ext: XcbExtension | _Pointer[XcbExtension]
+) -> QueryExtensionReply:
+    """Get extension data for the given extension.
+
+    Returns the extension data, which includes whether the extension is present
+    and its opcode information.
+    """
+    reply_p = LIB.xcb.xcb_get_extension_data(xcb_conn, ext)
+    return reply_p.contents
+
+
+def prefetch_extension_data(
+    xcb_conn: Connection | _Pointer[Connection], ext: XcbExtension | _Pointer[XcbExtension]
+) -> None:
+    """Prefetch extension data for the given extension.
+
+    This is a performance hint to XCB to fetch the extension data
+    asynchronously.
+    """
+    LIB.xcb.xcb_prefetch_extension_data(xcb_conn, ext)
+
+
+def generate_id(xcb_conn: Connection | _Pointer[Connection]) -> XID:
+    """Generate a new unique X resource ID.
+
+    Returns an XID that can be used to create new X resources.
+    """
+    return LIB.xcb.xcb_generate_id(xcb_conn)
+
+
+def get_setup(xcb_conn: Connection | _Pointer[Connection]) -> Setup:
+    """Get the connection setup information.
+
+    Returns the setup structure containing information about the X server,
+    including available screens, pixmap formats, etc.
+    """
+    setup_p = LIB.xcb.xcb_get_setup(xcb_conn)
+    return setup_p.contents
+
+
+# Connection management
+
+
 def initialize() -> None:
     LIB.initialize(callbacks=[xcbgen.initialize])
 
@@ -145,6 +192,12 @@ def connect(display: str | bytes | None = None) -> tuple[Connection, int]:
             msg += f"error code {conn_err}"
         raise XError(msg)
 
+    # Prefetch extension data for all extensions we support to populate XCB's internal cache.
+    prefetch_extension_data(conn_p, LIB.randr_id)
+    prefetch_extension_data(conn_p, LIB.render_id)
+    prefetch_extension_data(conn_p, LIB.shm_id)
+    prefetch_extension_data(conn_p, LIB.xfixes_id)
+
     return conn_p.contents, pref_screen_num.value