save_img(): little optimisation

BoboTiG · BoboTiG · commit cf85f2f31ac1 · 2015-04-12T18:21:05.000+02:00
diff --git a/mss.py b/mss.py
@@ -29,9 +29,8 @@
 
 from struct import pack
 from platform import system
-import os.path
+from zlib import compress, crc32
 import sys
-import zlib
 
 
 class ScreenshotError(Exception):
@@ -216,39 +215,37 @@ def save_img(self, data, width, height, output):
             http://inaps.org/journal/comment-fonctionne-le-png
         '''
 
-        self.debug('save_img')
-
-        to_take = (width * 3 + 3) & -4
-        padding = 0 if to_take % 8 == 0 else (to_take % 8) // 2
+        zcrc32 = crc32
+        zcompr = compress
+        len_sl = width * 3
         scanlines = b''.join(
-            [b'0' + data[(y * to_take):(y * to_take) + to_take - padding]
+            [b'0' + data[y * len_sl:y * len_sl + len_sl]
              for y in range(height)])
 
         magic = pack(b'>8B', 137, 80, 78, 71, 13, 10, 26, 10)
 
         # Header: size, marker, data, CRC32
         ihdr = [b'', b'IHDR', b'', b'']
         ihdr[2] = pack(b'>2I5B', width, height, 8, 2, 0, 0, 0)
-        ihdr[3] = pack(b'>I', zlib.crc32(b''.join(ihdr[1:3])) & 0xffffffff)
+        ihdr[3] = pack(b'>I', zcrc32(b''.join(ihdr[1:3])) & 0xffffffff)
         ihdr[0] = pack(b'>I', len(ihdr[2]))
 
         # Data: size, marker, data, CRC32
         idat = [b'', b'IDAT', b'', b'']
-        idat[2] = zlib.compress(scanlines, 9)
-        idat[3] = pack(b'>I', zlib.crc32(b''.join(idat[1:3])) & 0xffffffff)
+        idat[2] = zcompr(scanlines, 9)
+        idat[3] = pack(b'>I', zcrc32(b''.join(idat[1:3])) & 0xffffffff)
         idat[0] = pack(b'>I', len(idat[2]))
 
         # Footer: size, marker, None, CRC32
         iend = [b'', b'IEND', b'', b'']
-        iend[3] = pack(b'>I', zlib.crc32(iend[1]) & 0xffffffff)
+        iend[3] = pack(b'>I', zcrc32(iend[1]) & 0xffffffff)
         iend[0] = pack(b'>I', len(iend[2]))
 
         with open(output, 'wb') as fileh:
             fileh.write(
                 magic + b''.join(ihdr) + b''.join(idat) + b''.join(iend))
-        if not os.path.isfile(output):
-            msg = 'Impossible to write data to file "{}".'.format(output)
-            raise ScreenshotError(msg)
+            return
+        raise ScreenshotError('MSS: error writing data to "{}".'.format(output))
 
 
 class MSSMac(MSS):
