Mac: Fix memory leak (fix BoboTiG#24)

Mickaël Schoentgen · Mickaël Schoentgen · commit 9aca393d0633 · 2017-08-11T08:57:14.000+02:00
diff --git a/CHANGELOG b/CHANGELOG
@@ -8,6 +8,7 @@ dev     2017/xx/xx
       - tests: a lot of tests added for better coverage
       - MSS: possibility to use custom class to handle screen shot data
       - Mac: handle screenshot's width not divisible by 16 (fix #14, #19, #21)
+      - Mac: fix memory leak (fix #24)
       - Linux: handle bad display value
       - Windows: Take into account zoom factor for high-DPI displays (fix #20)
 
diff --git a/docs/source/developers.rst b/docs/source/developers.rst
@@ -31,7 +31,7 @@ How to test?
 
 Enable the developer mode::
 
-    $ sudo python setup.py develop
+    $ python -m pip install -e .
 
 Lauch the test suit::
 
diff --git a/mss/darwin.py b/mss/darwin.py
@@ -98,6 +98,7 @@ def _set_argtypes(self):
         self.core.CFDataGetBytePtr.argtypes = [ctypes.c_void_p]
         self.core.CFDataGetLength.argtypes = [ctypes.c_void_p]
         self.core.CGDataProviderRelease.argtypes = [ctypes.c_void_p]
+        self.core.CFRelease.argtypes = [ctypes.c_void_p]
 
     def _set_restypes(self):
         # type: () -> None
@@ -114,6 +115,7 @@ def _set_restypes(self):
         self.core.CFDataGetBytePtr.restype = ctypes.c_void_p
         self.core.CFDataGetLength.restype = ctypes.c_uint64
         self.core.CGDataProviderRelease.restype = ctypes.c_void_p
+        self.core.CFRelease.restype = ctypes.c_void_p
 
     @property
     def monitors(self):
@@ -183,12 +185,13 @@ def grab(self, monitor):
                 'CoreGraphics.CGWindowListCreateImage() failed.', locals())
 
         prov = self.core.CGImageGetDataProvider(image_ref)
-        data = self.core.CGDataProviderCopyData(prov)
-        data_ref = self.core.CFDataGetBytePtr(data)
-        buf_len = self.core.CFDataGetLength(data)
-        data = ctypes.cast(data_ref, ctypes.POINTER(ctypes.c_ubyte * buf_len))
-        data = data.contents
+        copy_data = self.core.CGDataProviderCopyData(prov)
+        data_ref = self.core.CFDataGetBytePtr(copy_data)
+        buf_len = self.core.CFDataGetLength(copy_data)
+        raw = ctypes.cast(data_ref, ctypes.POINTER(ctypes.c_ubyte * buf_len))
+        data = bytearray(raw.contents)
         self.core.CGDataProviderRelease(prov)
+        self.core.CFRelease(copy_data)
 
         if rounded_width != monitor['width']:
             data = self.resize(data, monitor)
diff --git a/tests/test_gnu_linux.py b/tests/test_gnu_linux.py
@@ -31,11 +31,13 @@ def test_factory_systems(monkeypatch):
     monkeypatch.setattr(platform, 'system', lambda: 'LINUX')
     sct = mss.mss()
     assert isinstance(sct, MSSBase)
+    monkeypatch.undo()
 
     # macOS
     monkeypatch.setattr(platform, 'system', lambda: 'Darwin')
     with pytest.raises(ScreenShotError) as exc:
         mss.mss()
+    monkeypatch.undo()
     if not PY3:
         error = exc.value[1]['self']
     else:
@@ -47,6 +49,7 @@ def test_factory_systems(monkeypatch):
     with pytest.raises(ValueError):
         # wintypes.py:19: ValueError: _type_ 'v' not supported
         mss.mss()
+    monkeypatch.undo()
 
 
 def test_implementation(monkeypatch, is_travis):
diff --git a/tests/test_implementation.py b/tests/test_implementation.py
@@ -63,7 +63,7 @@ def test_repr(sct):
     assert repr(img) == repr(ref)
 
 
-def test_factory_basics(monkeypatch):
+def test_factory(monkeypatch):
     # Current system
     with mss.mss() as sct:
         assert isinstance(sct, MSSBase)
@@ -72,13 +72,24 @@ def test_factory_basics(monkeypatch):
     monkeypatch.setattr(platform, 'system', lambda: 'Chuck Norris')
     with pytest.raises(ScreenShotError) as exc:
         mss.mss()
+    monkeypatch.undo()
+
     if not PY3:
         error = exc.value[0]
     else:
         error = exc.value.args[0]
     assert error == 'System not (yet?) implemented.'
 
 
-def test_python_call():
-    # __import__('mss.__main__')
-    pytest.skip('Dunno how to test mss/__main__.py.')
+def test_python_call(monkeypatch):
+    import mss.__main__
+    mss.__main__.main()
+
+    def raise_():
+        raise ScreenShotError()
+
+    pytest.skip('Not working for now.')
+    monkeypatch.setattr(mss.mss, '__init__', raise_)
+    with pytest.raises(ScreenShotError):
+        mss.mss()
+    monkeypatch.undo()
