Linux: added mouse support (partially fixes #55) (#232)

Original patch by @zorvios on #188.

Author: BoboTiG

CHANGELOG

@@ -3,6 +3,7 @@ History:
 <see Git checking messages for history>
 
 8.0.0   2023/0x/xx
+      - Linux: added mouse support (partially fixes #55)
       - Linux: removed get_error_details(), use the ScreenShotError details attribute instead
       - dev: removed pre-commit
       - tests: removed tox

CONTRIBUTORS

@@ -12,6 +12,9 @@ Alexander 'thehesiod' Mohr [https://github.com/thehesiod]
 Andreas Buhr [https://www.andreasbuhr.de]
     - Bugfix for multi-monitor detection
 
+Boutallaka 'zorvios' Yassir [https://github.com/zorvios]
+    - GNU/Linux: Mouse support
+
 bubulle [http://indexerror.net/user/bubulle]
     - Windows: efficiency of MSS.get_pixels()
 

docs/source/api.rst

@@ -25,13 +25,17 @@ GNU/Linux
 
 .. class:: MSS
 
-    .. method:: __init__([display=None])
+    .. method:: __init__([display=None, with_cursor=False])
 
         :type display: str or None
         :param display: The display to use.
+        :param with_cursor: Include the mouse cursor in screenshots.
 
         GNU/Linux initializations.
 
+            .. versionadded:: 8.0.0
+                `with_cursor` keyword argument.
+
     .. method:: grab(monitor)
 
         :rtype: :class:`~mss.base.ScreenShot`
@@ -76,6 +80,12 @@ Methods
 
     The parent's class for every OS implementation.
 
+    .. attribute:: compression_level
+
+        PNG compression level used when saving the screenshot data into a file (see :py:func:`zlib.compress()` for details).
+
+        .. versionadded:: 3.2.0
+
     .. method:: close()
 
         Clean-up method. Does nothing by default.

mss/base.py

@@ -18,11 +18,12 @@
 class MSSBase(metaclass=ABCMeta):
     """This class will be overloaded by a system specific one."""
 
-    __slots__ = {"_monitors", "cls_image", "compression_level"}
+    __slots__ = {"_monitors", "cls_image", "compression_level", "with_cursor"}
 
     def __init__(self) -> None:
         self.cls_image: Type[ScreenShot] = ScreenShot
         self.compression_level = 6
+        self.with_cursor = False
         self._monitors: Monitors = []
 
     def __enter__(self) -> "MSSBase":
@@ -35,6 +36,10 @@ def __exit__(self, *_: Any) -> None:
 
         self.close()
 
+    @abstractmethod
+    def _cursor_impl(self) -> Optional[ScreenShot]:
+        """Retrieve all cursor data. Pixels have to be RGB."""
+
     @abstractmethod
     def _grab_impl(self, monitor: Monitor) -> ScreenShot:
         """
@@ -73,7 +78,11 @@ def grab(self, monitor: Union[Monitor, Tuple[int, int, int, int]]) -> ScreenShot
             }
 
         with lock:
-            return self._grab_impl(monitor)
+            screenshot = self._grab_impl(monitor)
+            if self.with_cursor:
+                cursor = self._cursor_impl()
+                screenshot = self._merge(screenshot, cursor)  # type: ignore[arg-type]
+            return screenshot
 
     @property
     def monitors(self) -> Monitors:
@@ -174,6 +183,57 @@ def shot(self, **kwargs: Any) -> str:
         kwargs["mon"] = kwargs.get("mon", 1)
         return next(self.save(**kwargs))
 
+    @staticmethod
+    def _merge(screenshot: ScreenShot, cursor: ScreenShot) -> ScreenShot:
+        """Create composite image by blending screenshot and mouse cursor."""
+
+        # pylint: disable=too-many-locals,invalid-name
+
+        (cx, cy), (cw, ch) = cursor.pos, cursor.size
+        (x, y), (w, h) = screenshot.pos, screenshot.size
+
+        cx2, cy2 = cx + cw, cy + ch
+        x2, y2 = x + w, y + h
+
+        overlap = cx < x2 and cx2 > x and cy < y2 and cy2 > y
+        if not overlap:
+            return screenshot
+
+        screen_data = screenshot.raw
+        cursor_data = cursor.raw
+
+        cy, cy2 = (cy - y) * 4, (cy2 - y2) * 4
+        cx, cx2 = (cx - x) * 4, (cx2 - x2) * 4
+        start_count_y = -cy if cy < 0 else 0
+        start_count_x = -cx if cx < 0 else 0
+        stop_count_y = ch * 4 - max(cy2, 0)
+        stop_count_x = cw * 4 - max(cx2, 0)
+        rgb = range(3)
+
+        for count_y in range(start_count_y, stop_count_y, 4):
+            pos_s = (count_y + cy) * w + cx
+            pos_c = count_y * cw
+
+            for count_x in range(start_count_x, stop_count_x, 4):
+                spos = pos_s + count_x
+                cpos = pos_c + count_x
+                alpha = cursor_data[cpos + 3]
+
+                if not alpha:
+                    continue
+
+                if alpha == 255:
+                    screen_data[spos : spos + 3] = cursor_data[cpos : cpos + 3]
+                else:
+                    alpha = alpha / 255
+                    for i in rgb:
+                        screen_data[spos + i] = int(
+                            cursor_data[cpos + i] * alpha
+                            + screen_data[spos + i] * (1 - alpha)
+                        )
+
+        return screenshot
+
     @staticmethod
     def _cfactory(
         attr: Any,

mss/darwin.py

@@ -17,7 +17,7 @@
     c_void_p,
 )
 from platform import mac_ver
-from typing import Any, Type, Union
+from typing import Any, Optional, Type, Union
 
 from .base import MSSBase
 from .exception import ScreenShotError
@@ -232,3 +232,7 @@ def _grab_impl(self, monitor: Monitor) -> ScreenShot:
                 core.CFRelease(copy_data)
 
         return self.cls_image(data, monitor, size=Size(width, height))
+
+    def _cursor_impl(self) -> Optional[ScreenShot]:
+        """Retrieve all cursor data. Pixels have to be RGB."""
+        return None

mss/linux.py

@@ -15,12 +15,14 @@
     c_int,
     c_int32,
     c_long,
+    c_short,
     c_ubyte,
     c_uint,
     c_uint32,
     c_ulong,
     c_ushort,
     c_void_p,
+    cast,
 )
 from typing import Any, Dict, Optional, Tuple, Union
 
@@ -60,6 +62,26 @@ class Event(Structure):
     ]
 
 
+class XFixesCursorImage(Structure):
+    """
+    XFixes is an X Window System extension.
+    See /usr/include/X11/extensions/Xfixes.h
+    """
+
+    _fields_ = [
+        ("x", c_short),
+        ("y", c_short),
+        ("width", c_ushort),
+        ("height", c_ushort),
+        ("xhot", c_ushort),
+        ("yhot", c_ushort),
+        ("cursor_serial", c_ulong),
+        ("pixels", POINTER(c_ulong)),
+        ("atom", c_ulong),
+        ("name", c_char_p),
+    ]
+
+
 class XWindowAttributes(Structure):
     """Attributes for the specified window."""
 
@@ -211,6 +233,7 @@ def validate(retval: int, func: Any, args: Tuple[Any, Any]) -> Tuple[Any, Any]:
 CFUNCTIONS: CFunctions = {
     "XDefaultRootWindow": ("xlib", [POINTER(Display)], POINTER(XWindowAttributes)),
     "XDestroyImage": ("xlib", [POINTER(XImage)], c_void_p),
+    "XFixesGetCursorImage": ("xfixes", [POINTER(Display)], POINTER(XFixesCursorImage)),
     "XGetImage": (
         "xlib",
         [
@@ -269,15 +292,18 @@ class MSS(MSSBase):
     It uses intensively the Xlib and its Xrandr extension.
     """
 
-    __slots__ = {"drawable", "root", "xlib", "xrandr"}
+    __slots__ = {"drawable", "root", "xlib", "xrandr", "xfixes", "__with_cursor"}
 
     # A dict to maintain *display* values created by multiple threads.
     _display_dict: Dict[threading.Thread, int] = {}
 
-    def __init__(self, display: Optional[Union[bytes, str]] = None) -> None:
+    def __init__(
+        self, display: Optional[Union[bytes, str]] = None, with_cursor: bool = False
+    ) -> None:
         """GNU/Linux initialisations."""
 
         super().__init__()
+        self.with_cursor = with_cursor
 
         if not display:
             try:
@@ -306,6 +332,13 @@ def __init__(self, display: Optional[Union[bytes, str]] = None) -> None:
             raise ScreenShotError("No Xrandr extension found.")
         self.xrandr = ctypes.cdll.LoadLibrary(xrandr)
 
+        if self.with_cursor:
+            xfixes = ctypes.util.find_library("Xfixes")
+            if xfixes:
+                self.xfixes = ctypes.cdll.LoadLibrary(xfixes)
+            else:
+                self.with_cursor = False
+
         self._set_cfunctions()
 
         self.root = self.xlib.XDefaultRootWindow(self._get_display(display))
@@ -361,6 +394,7 @@ def _set_cfunctions(self) -> None:
         attrs = {
             "xlib": self.xlib,
             "xrandr": self.xrandr,
+            "xfixes": getattr(self, "xfixes", None),
         }
         for func, (attr, argtypes, restype) in CFUNCTIONS.items():
             with contextlib.suppress(AttributeError):
@@ -450,3 +484,32 @@ def _grab_impl(self, monitor: Monitor) -> ScreenShot:
             self.xlib.XDestroyImage(ximage)
 
         return self.cls_image(data, monitor)
+
+    def _cursor_impl(self) -> ScreenShot:
+        """Retrieve all cursor data. Pixels have to be RGB."""
+
+        # Read data of cursor/mouse-pointer
+        cursor_data = self.xfixes.XFixesGetCursorImage(self._get_display())
+        if not (cursor_data and cursor_data.contents):
+            raise ScreenShotError("Cannot read XFixesGetCursorImage()")
+
+        ximage: XFixesCursorImage = cursor_data.contents
+        monitor = {
+            "left": ximage.x - ximage.xhot,
+            "top": ximage.y - ximage.yhot,
+            "width": ximage.width,
+            "height": ximage.height,
+        }
+
+        raw_data = cast(
+            ximage.pixels, POINTER(c_ulong * monitor["height"] * monitor["width"])
+        )
+        raw = bytearray(raw_data.contents)
+
+        data = bytearray(monitor["height"] * monitor["width"] * 4)
+        data[3::4] = raw[3::8]
+        data[2::4] = raw[2::8]
+        data[1::4] = raw[1::8]
+        data[::4] = raw[::8]
+
+        return self.cls_image(data, monitor)

mss/tests/test_gnu_linux.py

@@ -132,3 +132,14 @@ def test_has_extension():
     with mss.mss(display=display) as sct:
         assert sct.has_extension("RANDR")
         assert not sct.has_extension("NOEXT")
+
+
+def test_with_cursor():
+    display = os.getenv("DISPLAY")
+    with mss.mss(display=display, with_cursor=True) as sct:
+        assert sct.xfixes
+        assert sct.with_cursor
+        sct.grab(sct.monitors[1])
+
+    # Not really sure how to test the cursor presence ...
+    # Also need to test when the cursor it outside of the screenshot

mss/windows.py

@@ -22,7 +22,7 @@
     UINT,
     WORD,
 )
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from .base import MSSBase
 from .exception import ScreenShotError
@@ -282,3 +282,7 @@ def _grab_impl(self, monitor: Monitor) -> ScreenShot:
             raise ScreenShotError("gdi32.GetDIBits() failed.")
 
         return self.cls_image(bytearray(self._data), monitor)
+
+    def _cursor_impl(self) -> Optional[ScreenShot]:
+        """Retrieve all cursor data. Pixels have to be RGB."""
+        return None