MSS: Add PNG compression level control

BoboTiG · web-flow · commit 5d3ab7c9e274 · 2018-01-20T07:52:47.000+01:00
diff --git a/CHANGELOG b/CHANGELOG
@@ -3,7 +3,7 @@ History:
 <see Git checking messages for history>
 
 dev     2018/xx/xx
-
+      - add PNG compression level control
 
 3.1.2   2018/01/05
       - removed support for Python 3.3
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -1,4 +1,16 @@
-3.1.2 (2018-01-xx)
+3.1.3 (2018-xx-xx)
+==================
+
+base.py
+-------
+- Added ``MSSBase.compression_level`` to control the PNG compression level
+
+tools.py
+--------
+- Changed signature of ``to_png(data, size, output=None)`` to ``to_png(data, size, level=6, output=None)``. ``level`` is the Zlib compression level.
+
+
+3.1.2 (2018-01-05)
 ==================
 
 tools.py
@@ -19,9 +31,9 @@ base.py
 
 darwin.py
 ---------
-- Add ``CGPoint.__repr__()``
-- Add ``CGRect.__repr__()``
-- Add ``CGSize.__repr__()``
+- Added ``CGPoint.__repr__()``
+- Added ``CGRect.__repr__()``
+- Added ``CGSize.__repr__()``
 - Removed ``get_infinity()`` function
 
 windows.py
@@ -36,7 +48,7 @@ windows.py
 base.py
 -------
 - Added the ``ScreenShot`` class containing data for a given screen shot (support the Numpy array interface [``ScreenShot.__array_interface__``])
-- Add ``shot()`` method to ``MSSBase``. It takes the same arguments as the ``save()`` method.
+- Added ``shot()`` method to ``MSSBase``. It takes the same arguments as the ``save()`` method.
 - Renamed ``get_pixels`` to ``grab``. It now returns a ``ScreenShot`` object.
 - Moved ``to_png`` method to ``tools.py``. It is now a simple function.
 - Removed ``enum_display_monitors()`` method. Use ``monitors`` property instead.
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -110,18 +110,24 @@ Methods
 
 .. module:: mss.tools
 
-.. method:: to_png(data, size, output=None)
+.. method:: to_png(data, size, level=6, output=None)
 
     :param bytes data: RGBRGB...RGB data.
     :param tuple size: The (width, height) pair.
+    :param int level: PNG compression level.
     :param str output: output's file name.
     :raises ScreenShotError: On error when writing ``data`` to ``output``.
+    :raises zlib.error: On bad compression ``level``.
 
     Dump data to the image file. Pure Python PNG implementation.
     If ``output`` is ``None``, create no file but return the whole PNG data.
 
     .. versionadded:: 3.0.0
 
+    .. versionadded:: 3.1.3
+
+        Added the ``level`` keyword argument to control the PNG compression level.
+
 
 Properties
 ==========
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -31,7 +31,7 @@
 # built documents.
 #
 # The short X.Y version.
-version = '3.1.2'
+version = '3.1.3'
 
 # The full version, including alpha/beta/rc tags.
 release = 'latest'
diff --git a/docs/source/examples.rst b/docs/source/examples.rst
@@ -56,6 +56,15 @@ This is an example that uses it, but also using percentage values:
 
 .. versionadded:: 3.1.0
 
+PNG Compression
+---------------
+
+You can tweak the PNG compression level (see :py:func:`zlib.compress()` for details)::
+
+    sct.compression_level = 2
+
+.. versionadded:: 3.1.3
+
 Advanced
 ========
 
diff --git a/docs/source/examples/from_pil_tuple.py b/docs/source/examples/from_pil_tuple.py
@@ -27,4 +27,4 @@
     im = sct.grab(bbox)
 
     # Save it!
-    mss.tools.to_png(im.rgb, im.size, 'screenshot.png')
+    mss.tools.to_png(im.rgb, im.size, output='screenshot.png')
diff --git a/docs/source/examples/part_of_screen.py b/docs/source/examples/part_of_screen.py
@@ -19,5 +19,5 @@
     sct_img = sct.grab(monitor)
 
     # Save to the picture file
-    mss.tools.to_png(sct_img.rgb, sct_img.size, output)
+    mss.tools.to_png(sct_img.rgb, sct_img.size, output=output)
     print(output)
diff --git a/docs/source/support.rst b/docs/source/support.rst
@@ -23,8 +23,8 @@ Tested successfully on Pypy 5.1.0 on Windows, but speed is terrible.
 Abandoned
 =========
 
-- Support for Python 2.6 was dropped on 2016-10-08: too old and introduced optimizations broke it.
-- Support for Python 3.0, 3.1, and 3.2 was dropped on 2016-10-08: there is no more tests facilities.
-- Support for Python 3.3 was dropped on 2017-12-05: there is no more tests facilities for the documentation.
+- 2016-10-08 Support for Python 2.6 was dropped: too old and introduced optimizations broke it.
+- 2016-10-08 Support for Python 3.0, 3.1, and 3.2 was dropped: there is no more tests facilities.
+- 2017-12-05 Support for Python 3.3 was dropped: there is no more tests facilities for the documentation.
 
 By the way, if you find the *force*, give it a try and tell us if you managed to make one of these versions working (a patch should be quite easy).
diff --git a/docs/source/where.rst b/docs/source/where.rst
@@ -13,6 +13,7 @@ AI, Computer Vison
 - `Europilot <https://github.com/jsistla/eu-pilot>`_, a self-driving algorithm using Euro Truck Simulator (ETS2);
 - `gym-mupen64plus <https://github.com/bzier/gym-mupen64plus>`_, an OpenAI Gym environment wrapper for the Mupen64Plus N64 emulator;
 - `Open Source Self Driving Car Initiative <https://github.com/OSSDC/OSSDC-VisionBasedACC>`_;
+- `PUBGIS <https://github.com/andrewzwicky/PUBGIS>`_, a map generator of your position throughout PUBG gameplay;
 - `Self-Driving-Car-3D-Simulator-With-CNN <https://github.com/sagar448/Self-Driving-Car-3D-Simulator-With-CNN>`_;
 
 Games
diff --git a/mss/__init__.py b/mss/__init__.py
@@ -13,7 +13,7 @@
 from .exception import ScreenShotError
 from .factory import mss
 
-__version__ = '3.1.2'
+__version__ = '3.1.3'
 __author__ = "Mickaël 'Tiger-222' Schoentgen"
 __copyright__ = """
     Copyright (c) 2013-2018, Mickaël 'Tiger-222' Schoentgen
diff --git a/mss/__main__.py b/mss/__main__.py
@@ -24,6 +24,9 @@ def main(args=None):
     cli_args.add_argument('-c', '--coordinates', default='', type=str,
                           help='the part of the screen to capture:'
                                ' top, left, width, height')
+    cli_args.add_argument('-l', '--level', default=6, type=int,
+                          choices=list(range(10)),
+                          help='the PNG compression level')
     cli_args.add_argument('-m', '--monitor', default=0, type=int,
                           help='the monitor to screen shot')
     cli_args.add_argument('-o', '--output', default='monitor-{mon}.png',
@@ -59,7 +62,10 @@ def main(args=None):
             if options.coordinates:
                 output = kwargs['output'].format(**kwargs['mon'])
                 sct_img = sct.grab(kwargs['mon'])
-                to_png(sct_img.rgb, sct_img.size, output)
+                to_png(sct_img.rgb,
+                       sct_img.size,
+                       level=options.level,
+                       output=output)
                 if not options.quiet:
                     print(os.path.realpath(output))
             else:
diff --git a/mss/base.py b/mss/base.py
@@ -15,6 +15,7 @@ class MSSBase(object):
     """ This class will be overloaded by a system specific one. """
 
     cls_image = ScreenShot  # type: object
+    compression_level = 6  # type: int
     _monitors = []  # type: List[Dict[str, int]]
 
     def __enter__(self):
@@ -106,7 +107,10 @@ def save(self, mon=0, output='monitor-{mon}.png', callback=None):
                 if callable(callback):
                     callback(fname)
                 sct = self.grab(monitor)
-                to_png(sct.rgb, sct.size, fname)
+                to_png(sct.rgb,
+                       sct.size,
+                       level=self.compression_level,
+                       output=fname)
                 yield fname
         else:
             # A screen shot of all monitors together or
@@ -121,7 +125,10 @@ def save(self, mon=0, output='monitor-{mon}.png', callback=None):
             if callable(callback):
                 callback(output)
             sct = self.grab(monitor)
-            to_png(sct.rgb, sct.size, output)
+            to_png(sct.rgb,
+                   sct.size,
+                   level=self.compression_level,
+                   output=output)
             yield output
 
     def shot(self, **kwargs):
diff --git a/mss/linux.py b/mss/linux.py
@@ -14,6 +14,10 @@
 __all__ = ('MSS',)
 
 
+PLAINMASK = 0x00ffffff
+ZPIXMAP = 2
+
+
 class Display(ctypes.Structure):
     """
     Structure that serves as the connection to the X server
@@ -289,7 +293,7 @@ def grab(self, monitor):
         ximage = self.xlib.XGetImage(self.display, root,
                                      monitor['left'], monitor['top'],
                                      monitor['width'], monitor['height'],
-                                     0x00ffffff, 2)  # ZPIXMAP
+                                     PLAINMASK, ZPIXMAP)
         if not ximage:
             raise ScreenShotError('xlib.XGetImage() failed.', locals())
 
diff --git a/mss/screenshot.py b/mss/screenshot.py
@@ -28,7 +28,8 @@ class ScreenShot(object):
 
     def __init__(self, data, monitor, size=None):
         # type: (bytearray, Dict[str, int], Any) -> None
-        #: Bytearray of the raw BGRA pixels retrieved by ctype
+
+        #: Bytearray of the raw BGRA pixels retrieved by ctypes
         #: OS independent implementations.
         self.raw = bytearray(data)  # type: bytearray
 
diff --git a/mss/tools.py b/mss/tools.py
@@ -8,14 +8,15 @@
 import zlib
 
 
-def to_png(data, size, output=None):
-    # type: (bytes, Tuple[int, int], Optional[str]) -> Union[None, bytes]
+def to_png(data, size, level=6, output=None):
+    # type: (bytes, Tuple[int, int], int, Optional[str]) -> Union[None, bytes]
     """
     Dump data to a PNG file.  If `output` is `None`, create no file but return
     the whole PNG data.
 
     :param bytes data: RGBRGB...RGB data.
     :param tuple size: The (width, height) pair.
+    :param int level: PNG compression level.
     :param str output: Output file name.
     """
 
@@ -35,7 +36,7 @@ def to_png(data, size, output=None):
     ihdr[0] = struct.pack('>I', len(ihdr[2]))
 
     # Data: size, marker, data, CRC32
-    idat = [b'', b'IDAT', zlib.compress(scanlines), b'']
+    idat = [b'', b'IDAT', zlib.compress(scanlines, level), b'']
     idat[3] = struct.pack('>I', zlib.crc32(b''.join(idat[1:3])) & 0xffffffff)
     idat[0] = struct.pack('>I', len(idat[2]))
 
diff --git a/tests/test_tools.py b/tests/test_tools.py
@@ -2,6 +2,9 @@
 
 import hashlib
 import os.path
+import zlib
+
+import pytest
 
 from mss.tools import to_png
 
@@ -11,6 +14,43 @@
 MD5SUM = '055e615b74167c9bdfea16a00539450c'
 
 
+def test_bad_compression_level(sct):
+    sct.compression_level = 42
+    try:
+        with pytest.raises(zlib.error):
+            sct.shot()
+    finally:
+        sct.compression_level = 6
+
+
+def test_compression_level(sct):
+    data = b'rgb' * WIDTH * HEIGHT
+    output = '{}x{}.png'.format(WIDTH, HEIGHT)
+
+    to_png(data, (WIDTH, HEIGHT), level=sct.compression_level, output=output)
+    with open(output, 'rb') as png:
+        assert hashlib.md5(png.read()).hexdigest() == MD5SUM
+
+
+@pytest.mark.parametrize('level, checksum', [
+    (0, 'f37123dbc08ed7406d933af11c42563e'),
+    (1, '7d5dcf2a2224445daf19d6d91cf31cb5'),
+    (2, 'bde05376cf51cf951e26c31c5f55e9d5'),
+    (3, '3d7e73c2a9c2d8842b363eeae8085919'),
+    (4, '9565a5caf89a9221459ee4e02b36bf6e'),
+    (5, '4d722e21e7d62fbf1e3154de7261fc67'),
+    (6, '055e615b74167c9bdfea16a00539450c'),
+    (7, '4d88d3f5923b6ef05b62031992294839'),
+    (8, '4d88d3f5923b6ef05b62031992294839'),
+    (9, '4d88d3f5923b6ef05b62031992294839'),
+])
+def test_compression_levels(level, checksum):
+    data = b'rgb' * WIDTH * HEIGHT
+    raw = to_png(data, (WIDTH, HEIGHT), level=level)
+    md5 = hashlib.md5(raw).hexdigest()
+    assert md5 == checksum
+
+
 def test_output_file():
     data = b'rgb' * WIDTH * HEIGHT
     output = '{}x{}.png'.format(WIDTH, HEIGHT)

Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@`
`31`	`31`	`# built documents.`
`32`	`32`	`#`
`33`	`33`	`# The short X.Y version.`
`34`		`-version = '3.1.2'`
	`34`	`+version = '3.1.3'`
`35`	`35`
`36`	`36`	`# The full version, including alpha/beta/rc tags.`
`37`	`37`	`release = 'latest'`