tools: force write of file when saving a PNG file

Mickaël Schoentgen · BoboTiG · commit cecafb24c455 · 2020-01-04T14:43:34.000+01:00
diff --git a/CHANGELOG b/CHANGELOG
@@ -4,6 +4,7 @@ History:
 
 5.0.1   2020/xx/xx
       - produce wheels for Python 3 only
+      - tools: force write of file when saving a PNG file
 
 5.0.0   2019/12/31
       - removed support for Python 2.7
diff --git a/mss/tools.py b/mss/tools.py
@@ -3,6 +3,7 @@
 Source: https://github.com/BoboTiG/python-mss
 """
 
+import os
 import struct
 import zlib
 from typing import TYPE_CHECKING
@@ -22,31 +23,35 @@ def to_png(data, size, level=6, output=None):
     :param int level: PNG compression level.
     :param str output: Output file name.
     """
+    # pylint: disable=too-many-locals
+
+    pack = struct.pack
+    crc32 = zlib.crc32
 
     width, height = size
     line = width * 3
-    png_filter = struct.pack(">B", 0)
+    png_filter = pack(">B", 0)
     scanlines = b"".join(
         [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)
+    magic = pack(">8B", 137, 80, 78, 71, 13, 10, 26, 10)
 
     # Header: size, marker, data, CRC32
     ihdr = [b"", b"IHDR", b"", b""]
-    ihdr[2] = struct.pack(">2I5B", width, height, 8, 2, 0, 0, 0)
-    ihdr[3] = struct.pack(">I", zlib.crc32(b"".join(ihdr[1:3])) & 0xFFFFFFFF)
-    ihdr[0] = struct.pack(">I", len(ihdr[2]))
+    ihdr[2] = pack(">2I5B", width, height, 8, 2, 0, 0, 0)
+    ihdr[3] = pack(">I", crc32(b"".join(ihdr[1:3])) & 0xFFFFFFFF)
+    ihdr[0] = pack(">I", len(ihdr[2]))
 
     # Data: size, marker, data, CRC32
     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]))
+    idat[3] = pack(">I", crc32(b"".join(idat[1:3])) & 0xFFFFFFFF)
+    idat[0] = pack(">I", len(idat[2]))
 
     # Footer: size, marker, None, CRC32
     iend = [b"", b"IEND", b"", b""]
-    iend[3] = struct.pack(">I", zlib.crc32(iend[1]) & 0xFFFFFFFF)
-    iend[0] = struct.pack(">I", len(iend[2]))
+    iend[3] = pack(">I", crc32(iend[1]) & 0xFFFFFFFF)
+    iend[0] = pack(">I", len(iend[2]))
 
     if not output:
         # Returns raw bytes of the whole PNG data
@@ -58,4 +63,8 @@ def to_png(data, size, level=6, output=None):
         fileh.write(b"".join(idat))
         fileh.write(b"".join(iend))
 
+        # Force write of file to disk
+        fileh.flush()
+        os.fsync(fileh.fileno())
+
     return None
