ArgbImage from arbitrary (but square) image data + guess rgb-mask types

icnsutil/ArgbImage.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 from typing import Union, Iterator, Optional
+from math import sqrt
 from . import IcnsType, PackBytes
 try:
     from PIL import Image
@@ -72,12 +73,23 @@ class ArgbImage:
         if is_argb or data[:4] == b'\x00\x00\x00\x00':
             data = data[4:]  # remove ARGB and it32 header
 
-        idat = PackBytes.unpack(data)
+        uncompressed_data = PackBytes.unpack(data)
-        iType = IcnsType.match_maxsize(len(idat), 'argb' if is_argb else 'rgb')
 
-        self.size = iType.size
+        self.channels = 4 if is_argb else 3
-        self.channels = iType.channels or 0
+        per_channel = len(uncompressed_data) // self.channels
-        self.a, self.r, self.g, self.b = iType.split_channels(idat)
+        w = sqrt(per_channel)
+        if w != int(w):
+            raise NotImplementedError(
+                'Could not determine square image size. Or unknown type.')
+        self.size = (int(w), int(w))
+        if self.channels == 3:
+            self.a = [255] * per_channel  # opaque alpha channel for rgb
+        else:
+            self.a = uncompressed_data[:per_channel]
+        i = 1 if is_argb else 0
+        self.r = uncompressed_data[(i + 0) * per_channel:(i + 1) * per_channel]
+        self.g = uncompressed_data[(i + 1) * per_channel:(i + 2) * per_channel]
+        self.b = uncompressed_data[(i + 2) * per_channel:(i + 3) * per_channel]
 
     def load_mask(self, *, file: Optional[str] = None,
                   data: Optional[bytes] = None) -> None:
@@ -99,15 +111,12 @@ class ArgbImage:
         return bytes(PackBytes.msb_stream(self.a, bits=bits))
 
     def rgb_data(self, *, compress: bool = True) -> bytes:
-        return b''.join(self._raw_rgb_channels(compress=compress))
+        return b''.join(PackBytes.pack(x) if compress else bytes(x)
+                        for x in (self.r, self.g, self.b))
 
     def argb_data(self, *, compress: bool = True) -> bytes:
-        return b'ARGB' + self.mask_data(compress=compress) + \
+        return b'ARGB' + self.mask_data(compress=compress) \
-            b''.join(self._raw_rgb_channels(compress=compress))
+                       + self.rgb_data(compress=compress)
-
-    def _raw_rgb_channels(self, *, compress: bool = True) -> Iterator[bytes]:
-        for x in (self.r, self.g, self.b):
-            yield (PackBytes.pack(x) if compress else bytes(x))
 
     def _load_png(self, fname: str) -> None:
         if not PIL_ENABLED:

icnsutil/IcnsType.py

@@ -54,20 +54,6 @@ class Media:
             return self.desc  # guaranteed to be icon, mask, or iconmask
         return self.types[-1]
 
-    def split_channels(self, uncompressed_data: List[int]) -> Iterator[
-            List[int]]:
-        if self.channels not in [3, 4]:
-            raise NotImplementedError('Only RGB and ARGB data supported.')
-        if len(uncompressed_data) != self.maxsize:
-            raise ValueError(
-                'Data does not match expected uncompressed length. '
-                '{} != {}'.format(len(uncompressed_data), self.maxsize))
-        per_channel = self.maxsize // self.channels
-        if self.channels == 3:
-            yield [255] * per_channel  # opaque alpha channel for rgb
-        for i in range(self.channels):
-            yield uncompressed_data[per_channel * i:per_channel * (i + 1)]
-
     def decompress(self, data: bytes, ext: Optional[str] = '-?-') -> Optional[
             List[int]]:
         ''' Returns None if media is not decompressable. '''
@@ -230,7 +216,7 @@ def get(key: Media.KeyT) -> Media:
 def match_maxsize(total: int, typ: str) -> Media:
     assert(typ == 'argb' or typ == 'rgb')
     ret = [x for x in _TYPES.values() if x.is_type(typ) and x.maxsize == total]
-    return ret[0]  # TODO: handle cases with multiple options? eg: is32 icp4
+    return _best_option(ret, typ)
 
 
 def guess(data: bytes, filename: Optional[str] = None) -> Media:
@@ -247,38 +233,64 @@ def guess(data: bytes, filename: Optional[str] = None) -> Media:
         if bname in _TYPES:
             return _TYPES[bname]
 
-    ext = RawData.determine_file_ext(data)
+    # Filter attributes
-    if not ext and filename and filename.endswith('.rgb'):
-        ext = 'rgb'
-
-    # Guess by image size and retina flag
-    size = RawData.determine_image_size(data, ext) if ext else None
-    retina = bname.lower().endswith('@2x') if filename else False
-    if size == (1024, 1024):
-        retina = True  # stupid double usage of ic10
-
-    # Icns specific names
     desc = None
-    if ext == 'icns' and filename:
+    size = None
-        for candidate in ['template', 'selected', 'dark']:
+    maxsize = None
-            if filename.endswith(candidate + '.icns'):
+    retina = False
-                desc = candidate
+
-                break
+    # Guess extension
+    ext = RawData.determine_file_ext(data)
+    if not ext and filename:
+        if filename.endswith('.rgb'):
+            ext = 'rgb'
+        elif filename.endswith('.mask'):
+            maxsize = len(data)
+            desc = 'mask'
+
+    # Guess image size
+    if ext:
+        size = RawData.determine_image_size(data, ext)
+
+    # if filename is set, then bname is also set (see above)
+    if filename:
+        # Guess retina flag
+        retina = bname.lower().endswith('@2x')
+        # Guess icns-specific type
+        if ext == 'icns':
+            for candidate in ['template', 'selected', 'dark']:
+                if bname.endswith(candidate):
+                    desc = candidate
+                    break
+
+    # stupid double usage of ic10, enforce retina flag
+    if size == (1024, 1024):
+        retina = True
 
     choices = []
     for x in _TYPES.values():
         if retina != x.retina:  # png + jp2
             continue
+        if desc and desc != x.desc:  # icns or rgb-mask
+            continue
         if ext:
             if size != x.size or not x.is_type(ext):
                 continue
-            if desc and desc != x.desc:  # icns only
-                continue
         else:  # not ext
             if x.ext_certain:
                 continue
+            if maxsize and x.maxsize and maxsize != x.maxsize:  # mask only
+                continue
         choices.append(x)
 
+    return _best_option(choices, ext)
+
+
+def _best_option(choices: List[Media], ext: Optional[str] = None) -> Media:
+    '''
+    Get most favorable media type.
+    If more than one option exists, choose based on order index of ext.
+    '''
     if len(choices) == 1:
         return choices[0]
     # Try get most favorable type (sort order of types)
@@ -294,7 +306,7 @@ def guess(data: bytes, filename: Optional[str] = None) -> Media:
                 best_choice.append(x)
         if len(best_choice) == 1:
             return best_choice[0]
-    # Else
+        choices = best_choice
-    raise CanNotDetermine(
+
-        'Could not determine type for file: "{}" – one of {}.'.format(
+    raise CanNotDetermine('Could not determine type – one of {}.'.format(
-            filename, [x.key for x in choices]))
+        [x.key for x in choices]))