Windows: refactored how internal handles are stored (#250)

Author: BoboTiG

CHANGELOG.md

@@ -4,9 +4,11 @@ See Git checking messages for full history.
 
 ## 8.0.4 (2023/xx/xx)
 - Linux: add failure handling to `XOpenDisplay()` call (#247)
+- Windows: refactored how internal handles are stored
+- Windows: removed side effects when leaving the context manager, resources are all freed
 - CI: run tests via xvfb-run on GitHub Actions (#248)
 - tests: Use `PyVirtualDisplay` instead of `xvfbwrapper` (#249)
-- :heart: contributors: @mgorny
+- :heart: contributors: @mgorny, @CTPaHHuK-HEbA
 
 ## 8.0.3 (2023/04/15)
 - added support for Python 3.12

src/mss/windows.py

@@ -4,8 +4,7 @@
 """
 import ctypes
 import sys
-import threading
-from ctypes import POINTER, WINFUNCTYPE, Structure, c_void_p
+from ctypes import POINTER, WINFUNCTYPE, Structure, c_int, c_void_p
 from ctypes.wintypes import (
     BOOL,
     DOUBLE,
@@ -22,7 +21,8 @@
     UINT,
     WORD,
 )
-from typing import Any, Dict, Optional
+from threading import local
+from typing import Any, Optional
 
 from .base import MSSBase
 from .exception import ScreenShotError
@@ -78,27 +78,22 @@ class BITMAPINFO(Structure):
     "BitBlt": ("gdi32", [HDC, INT, INT, INT, INT, HDC, INT, INT, DWORD], BOOL),
     "CreateCompatibleBitmap": ("gdi32", [HDC, INT, INT], HBITMAP),
     "CreateCompatibleDC": ("gdi32", [HDC], HDC),
+    "DeleteDC": ("gdi32", [HDC], HDC),
     "DeleteObject": ("gdi32", [HGDIOBJ], INT),
     "EnumDisplayMonitors": ("user32", [HDC, c_void_p, MONITORNUMPROC, LPARAM], BOOL),
     "GetDeviceCaps": ("gdi32", [HWND, INT], INT),
     "GetDIBits": ("gdi32", [HDC, HBITMAP, UINT, UINT, c_void_p, POINTER(BITMAPINFO), UINT], BOOL),
     "GetSystemMetrics": ("user32", [INT], INT),
     "GetWindowDC": ("user32", [HWND], HDC),
+    "ReleaseDC": ("user32", [HWND, HDC], c_int),
     "SelectObject": ("gdi32", [HDC, HGDIOBJ], HGDIOBJ),
 }
 
 
 class MSS(MSSBase):
     """Multiple ScreenShots implementation for Microsoft Windows."""
 
-    __slots__ = {"_bbox", "_bmi", "_data", "gdi32", "user32"}
-
-    # Class attributes instanced one time to prevent resource leaks.
-    bmp = None
-    memdc = None
-
-    # A dict to maintain *srcdc* values created by multiple threads.
-    _srcdc_dict: Dict[threading.Thread, int] = {}
+    __slots__ = {"gdi32", "user32", "_handles"}
 
     def __init__(self, /, **kwargs: Any) -> None:
         """Windows initialisations."""
@@ -110,12 +105,12 @@ def __init__(self, /, **kwargs: Any) -> None:
         self._set_cfunctions()
         self._set_dpi_awareness()
 
-        self._bbox = {"height": 0, "width": 0}
-        self._data: ctypes.Array[ctypes.c_char] = ctypes.create_string_buffer(0)
-
-        srcdc = self._get_srcdc()
-        if not MSS.memdc:
-            MSS.memdc = self.gdi32.CreateCompatibleDC(srcdc)
+        # Available thread-specific variables
+        self._handles = local()
+        self._handles.region_height_width = (0, 0)
+        self._handles.bmp = None
+        self._handles.srcdc = self.user32.GetWindowDC(0)
+        self._handles.memdc = self.gdi32.CreateCompatibleDC(self._handles.srcdc)
 
         bmi = BITMAPINFO()
         bmi.bmiHeader.biSize = ctypes.sizeof(BITMAPINFOHEADER)
@@ -124,7 +119,21 @@ def __init__(self, /, **kwargs: Any) -> None:
         bmi.bmiHeader.biCompression = 0  # 0 = BI_RGB (no compression)
         bmi.bmiHeader.biClrUsed = 0  # See grab.__doc__ [3]
         bmi.bmiHeader.biClrImportant = 0  # See grab.__doc__ [3]
-        self._bmi = bmi
+        self._handles.bmi = bmi
+
+    def close(self) -> None:
+        # Clean-up
+        if self._handles.bmp:
+            self.gdi32.DeleteObject(self._handles.bmp)
+            self._handles.bmp = None
+
+        if self._handles.memdc:
+            self.gdi32.DeleteDC(self._handles.memdc)
+            self._handles.memdc = None
+
+        if self._handles.srcdc:
+            self.user32.ReleaseDC(0, self._handles.srcdc)
+            self._handles.srcdc = None
 
     def _set_cfunctions(self) -> None:
         """Set all ctypes functions and attach them to attributes."""
@@ -149,26 +158,9 @@ def _set_dpi_awareness(self) -> None:
             #     These applications are not automatically scaled by the system.
             ctypes.windll.shcore.SetProcessDpiAwareness(2)
         elif (6, 0) <= version < (6, 3):
-            # Windows Vista, 7, 8 and Server 2012
+            # Windows Vista, 7, 8, and Server 2012
             self.user32.SetProcessDPIAware()
 
-    def _get_srcdc(self) -> int:
-        """
-        Retrieve a thread-safe HDC from GetWindowDC().
-        In multithreading, if the thread that creates *srcdc* is dead, *srcdc* will
-        no longer be valid to grab the screen. The *srcdc* attribute is replaced
-        with *_srcdc_dict* to maintain the *srcdc* values in multithreading.
-        Since the current thread and main thread are always alive, reuse their *srcdc* value first.
-        """
-        cur_thread, main_thread = threading.current_thread(), threading.main_thread()
-        current_srcdc = MSS._srcdc_dict.get(cur_thread) or MSS._srcdc_dict.get(main_thread)
-        if current_srcdc:
-            srcdc = current_srcdc
-        else:
-            srcdc = self.user32.GetWindowDC(0)
-            MSS._srcdc_dict[cur_thread] = srcdc
-        return srcdc
-
     def _monitors_impl(self) -> None:
         """Get positions of monitors. It will populate self._monitors."""
 
@@ -240,35 +232,26 @@ def _grab_impl(self, monitor: Monitor, /) -> ScreenShot:
         Thanks to http://stackoverflow.com/a/3688682
         """
 
-        srcdc, memdc = self._get_srcdc(), MSS.memdc
+        srcdc, memdc = self._handles.srcdc, self._handles.memdc
+        gdi = self.gdi32
         width, height = monitor["width"], monitor["height"]
 
-        if (self._bbox["height"], self._bbox["width"]) != (height, width):
-            self._bbox = monitor
-            self._bmi.bmiHeader.biWidth = width
-            self._bmi.bmiHeader.biHeight = -height  # Why minus? [1]
-            self._data = ctypes.create_string_buffer(width * height * 4)  # [2]
-            if MSS.bmp:
-                self.gdi32.DeleteObject(MSS.bmp)
-            MSS.bmp = self.gdi32.CreateCompatibleBitmap(srcdc, width, height)
-            self.gdi32.SelectObject(memdc, MSS.bmp)
-
-        self.gdi32.BitBlt(
-            memdc,
-            0,
-            0,
-            width,
-            height,
-            srcdc,
-            monitor["left"],
-            monitor["top"],
-            SRCCOPY | CAPTUREBLT,
-        )
-        bits = self.gdi32.GetDIBits(memdc, MSS.bmp, 0, height, self._data, self._bmi, DIB_RGB_COLORS)
+        if self._handles.region_height_width != (height, width):
+            self._handles.region_height_width = (height, width)
+            self._handles.bmi.bmiHeader.biWidth = width
+            self._handles.bmi.bmiHeader.biHeight = -height  # Why minus? [1]
+            self._handles.data = ctypes.create_string_buffer(width * height * 4)  # [2]
+            if self._handles.bmp:
+                gdi.DeleteObject(self._handles.bmp)
+            self._handles.bmp = gdi.CreateCompatibleBitmap(srcdc, width, height)
+            gdi.SelectObject(memdc, self._handles.bmp)
+
+        gdi.BitBlt(memdc, 0, 0, width, height, srcdc, monitor["left"], monitor["top"], SRCCOPY | CAPTUREBLT)
+        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.")
 
-        return self.cls_image(bytearray(self._data), monitor)
+        return self.cls_image(bytearray(self._handles.data), monitor)
 
     def _cursor_impl(self) -> Optional[ScreenShot]:
         """Retrieve all cursor data. Pixels have to be RGB."""

src/tests/test_windows.py

@@ -24,34 +24,46 @@ def test_implementation(monkeypatch):
 
 def test_region_caching():
     """The region to grab is cached, ensure this is well-done."""
-    from mss.windows import MSS
-
     with mss.mss() as sct:
-        # Reset the current BMP
-        if MSS.bmp:
-            sct.gdi32.DeleteObject(MSS.bmp)
-            MSS.bmp = None
-
         # Grab the area 1
         region1 = {"top": 0, "left": 0, "width": 200, "height": 200}
         sct.grab(region1)
-        bmp1 = id(MSS.bmp)
+        bmp1 = id(sct._handles.bmp)
 
         # Grab the area 2, the cached BMP is used
         # Same sizes but different positions
         region2 = {"top": 200, "left": 200, "width": 200, "height": 200}
         sct.grab(region2)
-        bmp2 = id(MSS.bmp)
+        bmp2 = id(sct._handles.bmp)
         assert bmp1 == bmp2
 
         # Grab the area 2 again, the cached BMP is used
         sct.grab(region2)
-        assert bmp2 == id(MSS.bmp)
+        assert bmp2 == id(sct._handles.bmp)
+
+
+def test_region_not_caching():
+    """The region to grab is not bad cached previous grab."""
+    grab1 = mss.mss()
+    grab2 = mss.mss()
+
+    region1 = {"top": 0, "left": 0, "width": 100, "height": 100}
+    region2 = {"top": 0, "left": 0, "width": 50, "height": 1}
+    grab1.grab(region1)
+    bmp1 = id(grab1._handles.bmp)
+    grab2.grab(region2)
+    bmp2 = id(grab2._handles.bmp)
+    assert bmp1 != bmp2
+
+    # Grab the area 1, is not bad cached BMP previous grab the area 2
+    grab1.grab(region1)
+    bmp1 = id(grab1._handles.bmp)
+    assert bmp1 != bmp2
 
 
 def run_child_thread(loops):
     for _ in range(loops):
-        with mss.mss() as sct:
+        with mss.mss() as sct:  # New sct for every loop
             sct.grab(sct.monitors[1])
 
 
@@ -66,3 +78,21 @@ def test_thread_safety():
     thread2.start()
     thread1.join()
     thread2.join()
+
+
+def run_child_thread_bbox(loops, bbox):
+    with mss.mss() as sct:  # One sct for all loops
+        for _ in range(loops):
+            sct.grab(bbox)
+
+
+def test_thread_safety_regions():
+    """Thread safety test for different regions
+    The following code will throw a ScreenShotError exception if thread-safety is not guaranted.
+    """
+    thread1 = threading.Thread(target=run_child_thread_bbox, args=(100, (0, 0, 100, 100)))
+    thread2 = threading.Thread(target=run_child_thread_bbox, args=(100, (0, 0, 50, 1)))
+    thread1.start()
+    thread2.start()
+    thread1.join()
+    thread2.join()