Windows: capture cursor

src/mss/windows.py

@@ -20,6 +20,8 @@
     RECT,
     UINT,
     WORD,
+    POINT,
+    HICON
 )
 from threading import local
 from typing import Any, Optional
@@ -38,7 +40,9 @@
 
 
 class BITMAPINFOHEADER(Structure):
-    """Information about the dimensions and color format of a DIB."""
+    """
+    Information about the dimensions and color format of a DIB.
+    """
 
     _fields_ = [
         ("biSize", DWORD),
@@ -51,7 +55,7 @@ class BITMAPINFOHEADER(Structure):
         ("biXPelsPerMeter", LONG),
         ("biYPelsPerMeter", LONG),
         ("biClrUsed", DWORD),
-        ("biClrImportant", DWORD),
+        ("biClrImportant", DWORD)
     ]
 
 
@@ -63,6 +67,33 @@ class BITMAPINFO(Structure):
     _fields_ = [("bmiHeader", BITMAPINFOHEADER), ("bmiColors", DWORD * 3)]
 
 
+class CURSORINFO(Structure):
+    """
+    Information about the cursor.
+    """
+
+    _fields_ = [
+        ("cbSize", DWORD),
+        ("flags", DWORD),
+        ("hCursor", HDC),
+        ("ptScreenPos", POINT)
+    ]
+
+
+class ICONINFO(Structure):
+    """
+    Information about an icon or cursor.
+    """
+
+    _fields_ = [
+        ("fIcon", BOOL),
+        ("xHotspot", DWORD),
+        ("yHotspot", DWORD),
+        ("hbmMask", HBITMAP),
+        ("hbmColor", HBITMAP)
+    ]
+
+
 MONITORNUMPROC = WINFUNCTYPE(INT, DWORD, DWORD, POINTER(RECT), DOUBLE)
 
 
@@ -80,9 +111,12 @@ class BITMAPINFO(Structure):
     "CreateCompatibleDC": ("gdi32", [HDC], HDC),
     "DeleteDC": ("gdi32", [HDC], HDC),
     "DeleteObject": ("gdi32", [HGDIOBJ], INT),
+    "DrawIcon": ("user32", [HDC, INT, INT, HICON], BOOL),
     "EnumDisplayMonitors": ("user32", [HDC, c_void_p, MONITORNUMPROC, LPARAM], BOOL),
+    "GetCursorInfo": ("user32", [POINTER(CURSORINFO)], BOOL),
     "GetDeviceCaps": ("gdi32", [HWND, INT], INT),
     "GetDIBits": ("gdi32", [HDC, HBITMAP, UINT, UINT, c_void_p, POINTER(BITMAPINFO), UINT], BOOL),
+    "GetIconInfo": ("user32", [HICON, POINTER(ICONINFO)], BOOL),
     "GetSystemMetrics": ("user32", [INT], INT),
     "GetWindowDC": ("user32", [HWND], HDC),
     "ReleaseDC": ("user32", [HWND, HDC], c_int),
@@ -121,6 +155,13 @@ def __init__(self, /, **kwargs: Any) -> None:
         bmi.bmiHeader.biClrImportant = 0  # See grab.__doc__ [3]
         self._handles.bmi = bmi
 
+        cursor_info = CURSORINFO()
+        cursor_info.cbSize = ctypes.sizeof(CURSORINFO)
+        self._handles.cursor_info = cursor_info
+
+        icon_info = ICONINFO()  # 'ii' felt uncomfortable
+        self._handles.icon_info = icon_info
+
     def close(self) -> None:
         # Clean-up
         if self._handles.bmp:
@@ -200,7 +241,7 @@ def _callback(monitor: int, data: HDC, rect: LPRECT, dc_: LPARAM) -> int:
         callback = MONITORNUMPROC(_callback)
         user32.EnumDisplayMonitors(0, 0, callback, 0)
 
-    def _grab_impl(self, monitor: Monitor, /) -> ScreenShot:
+    def _grab_impl(self, monitor: Monitor, /) -> Optional[ScreenShot]:
         """
         Retrieve all pixels from a monitor. Pixels have to be RGB.
 
@@ -254,5 +295,66 @@ def _grab_impl(self, monitor: Monitor, /) -> ScreenShot:
         return self.cls_image(bytearray(self._handles.data), monitor)
 
     def _cursor_impl(self) -> Optional[ScreenShot]:
-        """Retrieve all cursor data. Pixels have to be RGB."""
-        return None
+        """Retrieve all cursor data. Pixels have to be RGB.
+
+        [1] user32.DrawIcon(HDC(memdc), 0, 0, hcursor)
+        Sometimes the memdc value is greater than the 32 bit limit
+        and that results in
+        'ctypes.ArgumentError: argument 1: OverflowError: int too long to convert'
+        but casting it to HDC type seems to fix the issue.
+
+        [2] user32.GetIconInfo(hcursor, self._handles.icon_info)
+        GetIconInfo also returns the handle for mask bitmap and the handle for color bitmap
+        but the color bitmap handle is null in case of monochrome cursors.
+
+        [3] is_monochrome = self._handles.icon_info.hbmColor is None
+        The correct way to detect monochrome cursors seems to be a unique property of their
+        mask bitmap. The height of the mask bitmap of a monochrome cursor is twice its width
+        (https://learn.microsoft.com/en-us/windows/win32/api/winuser/ns-winuser-iconinfo)
+        But I cannot find the correct way of getting a bitmap's dimensions, therefore I
+        just eneded up checking if the color bitmap is null.
+
+        [4]
+        The data received using DrawIcon is in the format BGRA but in case of monochrome
+        cursors the alpha value of every pixel is 0 for some reason. Therefore, the alpha
+        value of every non black pixel has to be manually set to 255.
+        """
+        srcdc, memdc = self._handles.srcdc, self._handles.memdc
+        gdi, user32 = self.gdi32, self.user32
+        width, height = 32, 32
+        user32.GetCursorInfo(self._handles.cursor_info)
+        hcursor = self._handles.cursor_info.hCursor
+        pos_screen = self._handles.cursor_info.ptScreenPos
+
+        if self._handles.region_width_height != (width, height):
+            self._handles.region_width_height = (width, height)
+            self._handles.bmi.bmiHeader.biWidth = width
+            self._handles.bmi.bmiHeader.biHeight = -height
+            self._handles.data = ctypes.create_string_buffer(width * height * 4)
+            if self._handles.bmp:
+                gdi.DeleteObject(self._handles.bmp)
+            self._handles.bmp = gdi.CreateCompatibleBitmap(srcdc, width, height)
+            gdi.SelectObject(memdc, self._handles.bmp)
+
+        user32.DrawIcon(HDC(memdc), 0, 0, hcursor)  # Why HDC? [1]
+        bits = gdi.GetDIBits(memdc, self._handles.bmp, 0, height, self._handles.data, self._handles.bmi, DIB_RGB_COLORS)
+        if bits != height:
+            raise ScreenShotError("gdi32.GetDIBits() failed.")
+
+        user32.GetIconInfo(hcursor, self._handles.icon_info)  # [2]
+        is_monochrome = self._handles.icon_info.hbmColor is None  # [3]
+        ratio = ctypes.windll.shcore.GetScaleFactorForDevice(0) / 100
+        region = {
+            "left": round(pos_screen.x * ratio - self._handles.icon_info.xHotspot),
+            "top": round(pos_screen.y * ratio - self._handles.icon_info.yHotspot),
+            "width": 32,
+            "height": 32
+        }
+        data = bytearray(self._handles.data)
+        if is_monochrome:
+            for i in range(3, len(data), 4):  # [4]
+                if data[i-3:i] == b"\x00\x00\x00":
+                    data[i] = 0
+                else:
+                    data[i] = 255
+        return self.cls_image(data, region)