MSS: Possibility to get the whole PNG raw bytes

BoboTiG · web-flow · commit dd52984f78cd · 2018-01-02T21:01:36.000+01:00
diff --git a/CHANGELOG b/CHANGELOG
@@ -6,6 +6,7 @@ dev     2017/xx/xx
       - removed support for Python 3.3
       - CI: build the documentation
       - doc: improvements and fixes (fix #37)
+      - MSS: possibility to get the whole PNG raw bytes
 
 3.1.1   2017/11/27
       - MSS: add the 'mss' entry point
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -110,14 +110,15 @@ Methods
 
 .. module:: mss.tools
 
-.. method:: to_png(data, size, output)
+.. method:: to_png(data, size, output=None)
 
     :param bytes data: RGBRGB...RGB data.
     :param tuple size: The (width, height) pair.
     :param str output: output's file name.
     :raises ScreenShotError: On error when writing ``data`` to ``output``.
 
     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
 
diff --git a/mss/tools.py b/mss/tools.py
@@ -8,10 +8,11 @@
 import zlib
 
 
-def to_png(data, size, output):
-    # type: (bytes, Tuple[int, int], str) -> None
+def to_png(data, size, output=None):
+    # type: (bytes, Tuple[int, int], Optional[str]) -> Union[None, bytes]
     """
-    Dump data to a PNG file.
+    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.
@@ -43,6 +44,10 @@ def to_png(data, size, output):
     iend[3] = struct.pack('>I', zlib.crc32(iend[1]) & 0xffffffff)
     iend[0] = struct.pack('>I', len(iend[2]))
 
+    if not output:
+        # Returns raw bytes of the whole PNG data
+        return magic + b''.join(ihdr + idat + iend)
+
     with open(output, 'wb') as fileh:
         fileh.write(magic)
         fileh.write(b''.join(ihdr))
diff --git a/setup.py b/setup.py
@@ -33,7 +33,6 @@
     'Programming Language :: Python :: 2',
     'Programming Language :: Python :: 2.7',
     'Programming Language :: Python :: 3',
-    'Programming Language :: Python :: 3.3',
     'Programming Language :: Python :: 3.4',
     'Programming Language :: Python :: 3.5',
     'Programming Language :: Python :: 3.6',
diff --git a/tests/test_tools.py b/tests/test_tools.py
@@ -0,0 +1,27 @@
+# coding: utf-8
+
+import hashlib
+import os.path
+
+from mss.tools import to_png
+
+
+WIDTH = 10
+HEIGHT = 10
+MD5SUM = '055e615b74167c9bdfea16a00539450c'
+
+
+def test_output_file():
+    data = b'rgb' * WIDTH * HEIGHT
+    output = '{}x{}.png'.format(WIDTH, HEIGHT)
+    to_png(data, (WIDTH, HEIGHT), output=output)
+
+    assert os.path.isfile(output)
+    with open(output, 'rb') as png:
+        assert hashlib.md5(png.read()).hexdigest() == MD5SUM
+
+
+def test_output_raw_bytes():
+    data = b'rgb' * WIDTH * HEIGHT
+    raw = to_png(data, (WIDTH, HEIGHT))
+    assert hashlib.md5(raw).hexdigest() == MD5SUM
