MSS: ensure calls without context manager will not leak resources or document them (BoboTiG#72, BoboTiG#85)

BoboTiG · BoboTiG · commit 1c914df043de · 2018-12-27T22:59:05.000+01:00
- Added test_leaks.py to ensure fixes are good and prevent regressions
diff --git a/CHANGELOG b/CHANGELOG
@@ -6,6 +6,7 @@ History:
       - MSS: renamed MSSBase to MSSMixin in base.py
       - Linux: ensure resources are freed in grab()
       - Windows: avoid unecessary class attributes
+      - MSS: ensure calls without context manager will not leak resources or document them (fixes #72 and #85)
       - MSS: fix Flake8 C408: Unnecessary dict call - rewrite as a literal, in exceptions.py
       - MSS: fix Flake8 I100: Import statements are in the wrong order
       - MSS: fix Flake8 I201: Missing newline before sections or imports
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -3,7 +3,15 @@
 
 base.py
 -------
-- Renamed ``MSSBase`` class to ``MSSMixin``.
+- Renamed ``MSSBase`` class to ``MSSMixin``
+
+linux.py
+--------
+- Renamed ``__del__()`` method to ``close()``
+
+windows.py
+----------
+- Renamed ``__exit__()`` method to ``close()``
 
 
 3.3.0 (2018-09-04)
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -19,6 +19,12 @@ GNU/Linux
 
         GNU/Linux initializations.
 
+    .. method:: close()
+
+        Disconnect from the X server.
+
+        .. versionadded:: 4.0.0
+
     .. method:: grab(monitor)
 
         :rtype: :class:`mss.base.ScreenShot`
@@ -41,6 +47,21 @@ GNU/Linux
 
     .. versionadded:: 3.3.0
 
+
+Windows
+-------
+
+.. module:: mss.windows
+
+.. class:: MSS
+
+    .. method:: close()
+
+        Close GDI handles and free DCs.
+
+        .. versionadded:: 4.0.0
+
+
 Methods
 =======
 
@@ -50,7 +71,11 @@ Methods
 
     The parent's class for every OS implementation.
 
-    .. versionchanged:: 4.0.0
+    .. method:: close()
+
+        Clean-up method. Does nothing by default.
+
+        .. versionadded:: 4.0.0
 
     .. method:: grab(region)
 
@@ -148,6 +173,8 @@ Methods
 Properties
 ==========
 
+.. module:: mss.base
+
 .. class:: MSSMixin
 
     .. attribute:: monitors
diff --git a/docs/source/examples/fps.py b/docs/source/examples/fps.py
@@ -61,6 +61,7 @@ def screen_record_efficient():
             cv2.destroyAllWindows()
             break
 
+    sct.close()
     return fps
 
 
diff --git a/docs/source/usage.rst b/docs/source/usage.rst
@@ -30,35 +30,38 @@ So the module can be used as simply as::
     with mss() as sct:
         # ...
 
-Or::
+.. note::
 
-    sct = mss()
+    On GNU/Linux and Windows, if you are using this kind of code::
 
+        sct = mss()
+        sct.shot()
+        sct.grab()
+        # or any attribute/method of sct
 
-GNU/Linux
----------
+    Then you will have to **manually** call :meth:`mss.linux.MSSMixin.close()` to free resources.
 
-On GNU/Linux, you can specify which display to use (useful for distant screenshots via SSH)::
+.. warning::
 
-    with mss(display=':0.0') as sct:
-        # ...
+    This code is **highly** unadvised as there will be **resources leaks** without possibility to do something to clean them::
 
-A more specific example to only target GNU/Linux:
-
-.. literalinclude:: examples/linux_display_keyword.py
-    :lines: 9-
+        mss().shot()
+        mss().grab()
+        # or any mss().xxx or mss().xxx()
 
 
 Intensive Use
 =============
 
-If you plan to integrate MSS inside your own module or software, pay attention to using it wisely. This is a bad usage::
+If you plan to integrate MSS inside your own module or software, pay attention to using it wisely.
+
+This is a bad usage::
 
     for _ in range(100):
         with mss() as sct:
             sct.shot()
 
-This is a better usage, memory efficient::
+This is a much better usage, memory efficient::
 
     with mss() as sct:
         for _ in range(100):
@@ -67,14 +70,28 @@ This is a better usage, memory efficient::
 Also, it is a good thing to save the MSS instance inside an attribute of you class and calling it when needed.
 
 
+GNU/Linux
+---------
+
+On GNU/Linux, you can specify which display to use (useful for distant screenshots via SSH)::
+
+    with mss(display=":0.0") as sct:
+        # ...
+
+A more specific example to only target GNU/Linux:
+
+.. literalinclude:: examples/linux_display_keyword.py
+    :lines: 9-
+
+
 Command Line
 ============
 
-You can use ``mss`` via the CLI:
+You can use ``mss`` via the CLI::
 
     mss --help
 
-Or via direct call from Python:
+Or via direct call from Python::
 
     python -m mss --help
 
diff --git a/mss/base.py b/mss/base.py
@@ -27,6 +27,14 @@ def __exit__(self, *_):
         # type: (*str) -> None
         """ For the cool call `with MSS() as mss:`. """
 
+        self.close()
+
+    def close(self):
+        # type: () -> None
+        """ Clean-up. """
+
+        pass
+
     def grab(self, monitor):
         # type: (Dict[str, int]) -> ScreenShot
         """
diff --git a/mss/linux.py b/mss/linux.py
@@ -146,16 +146,6 @@ class MSS(MSSMixin):
 
     last_error = None
 
-    def __del__(self):
-        # type: () -> None
-        """ Clean-up. """
-
-        # Disconnect from the X server
-        try:
-            self.xlib.XCloseDisplay(self.display)
-        except AttributeError:
-            pass
-
     def __init__(self, display=None):
         # type: (bytes) -> None
         """ GNU/Linux initialisations. """
@@ -341,6 +331,20 @@ def validate(retval, func, args):
         if errcheck:
             meth.errcheck = validate
 
+    def close(self):
+        # type: () -> None
+        """
+        Disconnect from the X server to prevent:
+            Maximum number of clients reached. Segmentation fault (core dumped)
+        """
+
+        try:
+            self.xlib.XCloseDisplay(self.display)
+            # Delete the attribute to prevent interpreter crash if called twice
+            del self.display
+        except AttributeError:
+            pass
+
     @property
     def monitors(self):
         # type: () -> List[Dict[str, int]]
diff --git a/mss/tools.py b/mss/tools.py
@@ -24,7 +24,7 @@ def to_png(data, size, level=6, output=None):
     line = width * 3
     png_filter = struct.pack(">B", 0)
     scanlines = b"".join(
-        [png_filter + data[y * line:y * line + line] for y in range(height)]
+        [png_filter + data[y * line : y * line + line] for y in range(height)]
     )
 
     magic = struct.pack(">8B", 137, 80, 78, 71, 13, 10, 26, 10)
diff --git a/mss/windows.py b/mss/windows.py
@@ -93,16 +93,27 @@ def __init__(self):
         bmi.bmiHeader.biClrImportant = 0  # See grab.__doc__ [3]
         self._bmi = bmi
 
-    def __exit__(self, *args):
-        # type: (*str) -> None
-        """ Cleanup. """
+    def close(self):
+        # type: () -> None
+        """ Close GDI handles and free DCs. """
 
-        for attribute in {"_bmp", "_memdc", "_srcdc"}:
-            attr = getattr(self, attribute, None)
-            if attr:
-                ctypes.windll.gdi32.DeleteObject(attr)
+        try:
+            ctypes.windll.gdi32.DeleteObject(self._bmp)
+            del self._bmp
+        except AttributeError:
+            pass
+
+        try:
+            ctypes.windll.gdi32.DeleteDC(self._memdc)
+            del self._memdc
+        except (OSError, AttributeError):
+            pass
 
-        super(MSS, self).__exit__(*args)
+        try:
+            ctypes.windll.user32.ReleaseDC(0, self._srcdc)
+            del self._srcdc
+        except AttributeError:
+            pass
 
     @property
     def monitors(self):
@@ -235,6 +246,10 @@ def set_argtypes(callback):
         ctypes.wintypes.LPARAM,
     ]
     ctypes.windll.user32.GetWindowDC.argtypes = [ctypes.wintypes.HWND]
+    ctypes.windll.user32.ReleaseDC.argtypes = [
+        ctypes.wintypes.HWND,
+        ctypes.wintypes.HGDIOBJ,
+    ]
     ctypes.windll.gdi32.GetDeviceCaps.argtypes = [
         ctypes.wintypes.HWND,
         ctypes.wintypes.INT,
@@ -279,6 +294,7 @@ def set_restypes():
     ctypes.windll.user32.GetSystemMetrics.restype = ctypes.wintypes.INT
     ctypes.windll.user32.EnumDisplayMonitors.restype = ctypes.wintypes.BOOL
     ctypes.windll.user32.GetWindowDC.restype = ctypes.wintypes.HDC
+    ctypes.windll.user32.ReleaseDC.restype = ctypes.wintypes.INT
     ctypes.windll.gdi32.GetDeviceCaps.restype = ctypes.wintypes.INT
     ctypes.windll.gdi32.CreateCompatibleDC.restype = ctypes.wintypes.HDC
     ctypes.windll.gdi32.CreateCompatibleBitmap.restype = ctypes.wintypes.HBITMAP
diff --git a/tests/leaks.py b/tests/leaks.py
diff --git a/tests/test_leaks.py b/tests/test_leaks.py

Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,7 @@ def to_png(data, size, level=6, output=None):`
`24`	`24`	`line = width * 3`
`25`	`25`	`png_filter = struct.pack(">B", 0)`
`26`	`26`	`scanlines = b"".join(`
`27`		`- [png_filter + data[y * line:y * line + line] for y in range(height)]`
	`27`	`+ [png_filter + data[y * line : y * line + line] for y in range(height)]`
`28`	`28`	`)`
`29`	`29`
`30`	`30`	`magic = struct.pack(">8B", 137, 80, 78, 71, 13, 10, 26, 10)`