230 lines
10 KiB
Python
230 lines
10 KiB
Python
"""XCB backend using MIT-SHM XShmGetImage with automatic fallback.
|
|
|
|
This is the fastest Linux backend available, and will work in most common
|
|
cases. However, it will not work over remote X connections, such as over ssh.
|
|
|
|
This implementation prefers shared-memory captures for performance and will
|
|
fall back to XGetImage when the MIT-SHM extension is unavailable or fails at
|
|
runtime. The fallback reason is exposed via ``performance_status`` to aid
|
|
debugging.
|
|
|
|
.. versionadded:: 10.2.0
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import enum
|
|
import os
|
|
from mmap import PROT_READ, mmap # type: ignore[attr-defined]
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from mss.exception import ScreenShotError
|
|
from mss.linux import xcb
|
|
from mss.linux.base import ALL_PLANES, MSSImplXCBBase
|
|
from mss.linux.xcbhelpers import LIB, XProtoError
|
|
|
|
if TYPE_CHECKING:
|
|
from mss.models import Monitor
|
|
|
|
__all__ = ()
|
|
|
|
|
|
class ShmStatus(enum.Enum):
|
|
"""Availability of the MIT-SHM extension for this backend."""
|
|
|
|
UNKNOWN = enum.auto() # Constructor says SHM *should* work, but we haven't seen a real GetImage succeed yet.
|
|
AVAILABLE = enum.auto() # We've successfully used XShmGetImage at least once.
|
|
UNAVAILABLE = enum.auto() # We know SHM GetImage is unusable; always use XGetImage.
|
|
|
|
|
|
class MSSImplXShmGetImage(MSSImplXCBBase):
|
|
"""XCB backend using XShmGetImage with an automatic XGetImage fallback.
|
|
|
|
.. seealso::
|
|
:py:class:`mss.linux.base.MSSImplXCBBase`
|
|
Lists constructor parameters.
|
|
"""
|
|
|
|
def __init__(self, *, display: str | bytes | None = None, with_cursor: bool = False) -> None:
|
|
super().__init__(display=display, with_cursor=with_cursor)
|
|
|
|
# These are the objects we need to clean up when we shut down. They are created in _setup_shm.
|
|
self._memfd: int | None = None
|
|
self._buf: mmap | None = None
|
|
self._shmseg: xcb.ShmSeg | None = None
|
|
|
|
# Rather than trying to track the shm_status, we may be able to raise an exception in __init__ if XShmGetImage
|
|
# isn't available. The factory in linux/__init__.py could then catch that and switch to XGetImage.
|
|
# The conditions under which the attach will succeed but the xcb_shm_get_image will fail are extremely
|
|
# rare, and I haven't yet found any that also will work with xcb_get_image.
|
|
#: Whether we can use the MIT-SHM extensions for this connection.
|
|
#: This will not be ``AVAILABLE`` until at least one capture has succeeded.
|
|
#: It may be set to ``UNAVAILABLE`` sooner.
|
|
self.shm_status: ShmStatus = self._setup_shm()
|
|
|
|
def _shm_report_issue(self, msg: str, *args: Any) -> None:
|
|
"""Debugging hook for troubleshooting MIT-SHM issues.
|
|
|
|
This will be called whenever MIT-SHM is disabled. The optional
|
|
arguments are not well-defined; exceptions are common.
|
|
"""
|
|
full_msg = msg
|
|
if args:
|
|
full_msg += " | " + ", ".join(str(arg) for arg in args)
|
|
self.performance_status.append(full_msg)
|
|
|
|
def _setup_shm(self) -> ShmStatus: # noqa: PLR0911
|
|
assert self.conn is not None # noqa: S101
|
|
|
|
try:
|
|
shm_ext_data = xcb.get_extension_data(self.conn, LIB.shm_id)
|
|
if not shm_ext_data.present:
|
|
self._shm_report_issue("MIT-SHM extension not present")
|
|
return ShmStatus.UNAVAILABLE
|
|
|
|
# We use the FD-based version of ShmGetImage, so we require the extension to be at least 1.2.
|
|
shm_version_data = xcb.shm_query_version(self.conn)
|
|
shm_version = (shm_version_data.major_version, shm_version_data.minor_version)
|
|
if shm_version < (1, 2):
|
|
self._shm_report_issue("MIT-SHM version too old", shm_version)
|
|
return ShmStatus.UNAVAILABLE
|
|
|
|
# We allocate something large enough for the root, so we don't have to reallocate each time the window is
|
|
# resized.
|
|
self._bufsize = self.pref_screen.width_in_pixels * self.pref_screen.height_in_pixels * 4
|
|
|
|
if not hasattr(os, "memfd_create"):
|
|
self._shm_report_issue("os.memfd_create not available")
|
|
return ShmStatus.UNAVAILABLE
|
|
try:
|
|
self._memfd = os.memfd_create("mss-shm-buf", flags=os.MFD_CLOEXEC) # type: ignore[attr-defined]
|
|
except OSError as e:
|
|
return self._shm_unavailable("memfd_create failed", e)
|
|
os.ftruncate(self._memfd, self._bufsize)
|
|
|
|
try:
|
|
self._buf = mmap(self._memfd, self._bufsize, prot=PROT_READ) # type: ignore[call-arg]
|
|
except OSError as e:
|
|
return self._shm_unavailable("mmap failed", e)
|
|
self._shmseg = xcb.ShmSeg(xcb.generate_id(self.conn).value)
|
|
try:
|
|
# This will normally be what raises an exception if you're on a remote connection.
|
|
# XCB will close _memfd, on success or on failure.
|
|
try:
|
|
xcb.shm_attach_fd(self.conn, self._shmseg, self._memfd, read_only=False)
|
|
finally:
|
|
self._memfd = None
|
|
except xcb.XError as e:
|
|
return self._shm_unavailable("Cannot attach MIT-SHM segment", e)
|
|
except Exception:
|
|
self._shutdown_shm()
|
|
raise
|
|
|
|
return ShmStatus.UNKNOWN
|
|
|
|
def _shm_unavailable(self, msg: str, exc: Exception) -> ShmStatus:
|
|
self._shm_report_issue(msg, exc)
|
|
self._shutdown_shm()
|
|
return ShmStatus.UNAVAILABLE
|
|
|
|
def close(self) -> None:
|
|
self._shutdown_shm()
|
|
super().close()
|
|
|
|
def _shutdown_shm(self) -> None:
|
|
# It would be nice to also try to tell the server to detach the shmseg, but we might be in an error path
|
|
# and don't know if that's possible. It's not like we'll leak a lot of them on the same connection anyway.
|
|
# This can be called in the path of partial initialization.
|
|
if self._buf is not None:
|
|
self._buf.close()
|
|
self._buf = None
|
|
if self._memfd is not None:
|
|
os.close(self._memfd)
|
|
self._memfd = None
|
|
|
|
def _grab_xshmgetimage(self, monitor: Monitor) -> bytearray:
|
|
if self.conn is None:
|
|
msg = "Cannot take screenshot while the connection is closed"
|
|
raise ScreenShotError(msg)
|
|
assert self._buf is not None # noqa: S101
|
|
assert self._shmseg is not None # noqa: S101
|
|
|
|
required_size = monitor["width"] * monitor["height"] * 4
|
|
if required_size > self._bufsize:
|
|
# This is temporary. The permanent fix will depend on how
|
|
# issue https://github.com/BoboTiG/python-mss/issues/432 is resolved.
|
|
msg = (
|
|
"Requested capture size exceeds the allocated buffer. If you have resized the screen, "
|
|
"please recreate your MSS object."
|
|
)
|
|
raise ScreenShotError(msg)
|
|
|
|
img_reply = xcb.shm_get_image(
|
|
self.conn,
|
|
self.drawable,
|
|
monitor["left"],
|
|
monitor["top"],
|
|
monitor["width"],
|
|
monitor["height"],
|
|
ALL_PLANES,
|
|
xcb.ImageFormat.ZPixmap,
|
|
self._shmseg,
|
|
0,
|
|
)
|
|
|
|
if img_reply.depth != self.drawable_depth or img_reply.visual != 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.value)}, "
|
|
f"got {img_reply.depth},{hex(img_reply.visual.value)}"
|
|
)
|
|
raise ScreenShotError(msg)
|
|
|
|
# Snapshot the buffer into new bytearray.
|
|
new_size = monitor["width"] * monitor["height"] * 4
|
|
# Slicing the memoryview creates a new memoryview that points to the relevant subregion. Making this and then
|
|
# copying it into a fresh bytearray is much faster than slicing the mmap object. Make sure we don't hold an
|
|
# open memoryview if an exception happens, since that will prevent us from closing self._buf during the stack
|
|
# unwind.
|
|
with memoryview(self._buf) as img_mv:
|
|
return bytearray(img_mv[:new_size])
|
|
|
|
def grab(self, monitor: Monitor) -> bytearray:
|
|
"""Retrieve all pixels from a monitor. Pixels have to be RGBX."""
|
|
if self.shm_status == ShmStatus.UNAVAILABLE:
|
|
return super()._grab_xgetimage(monitor)
|
|
|
|
# The usual path is just the next few lines.
|
|
try:
|
|
rv = self._grab_xshmgetimage(monitor)
|
|
if self.shm_status != ShmStatus.AVAILABLE:
|
|
self.shm_status = ShmStatus.AVAILABLE
|
|
self.performance_status.append("MIT-SHM is working correctly.")
|
|
except XProtoError as e:
|
|
if self.shm_status != ShmStatus.UNKNOWN:
|
|
# We know XShmGetImage works, because it worked earlier. Reraise the error.
|
|
raise
|
|
|
|
# Should we engage the fallback path? In almost all cases, if XShmGetImage failed at this stage (after
|
|
# all our testing in __init__), XGetImage will also fail. This could mean that the user sent an
|
|
# out-of-bounds request. In more exotic situations, some rare X servers disallow screen capture
|
|
# altogether: security-hardened servers, for instance, or some XPrint servers. But let's make sure, by
|
|
# testing the same request through XGetImage.
|
|
try:
|
|
rv = super()._grab_xgetimage(monitor)
|
|
except XProtoError: # noqa: TRY203
|
|
# The XGetImage also failed, so we don't know anything about whether XShmGetImage is usable. Maybe
|
|
# the user sent an out-of-bounds request. Maybe it's a security-hardened server. We're not sure what
|
|
# the problem is. So, if XGetImage failed, we re-raise that error (the one from XShmGetImage will be
|
|
# attached as __context__), but we won't update the shm_status yet. (Technically, our except:raise
|
|
# clause here is redundant; it's just for clarity, to hold this comment.)
|
|
raise
|
|
|
|
# Using XShmGetImage failed, and using XGetImage worked. Use XGetImage in the future.
|
|
self._shm_report_issue("MIT-SHM GetImage failed", e)
|
|
self.shm_status = ShmStatus.UNAVAILABLE
|
|
self._shutdown_shm()
|
|
|
|
return rv
|