From 8ca46c7f40200d60318127601c71889e1249a3ad Mon Sep 17 00:00:00 2001 From: relikd Date: Mon, 11 Oct 2021 23:56:35 +0200 Subject: [PATCH] ArgbImage from arbitrary (but square) image data + guess rgb-mask types --- icnsutil/ArgbImage.py | 33 +++++++++++------ icnsutil/IcnsType.py | 86 ++++++++++++++++++++++++------------------- 2 files changed, 70 insertions(+), 49 deletions(-) diff --git a/icnsutil/ArgbImage.py b/icnsutil/ArgbImage.py index 8dcb025..c5cc85e 100644 --- a/icnsutil/ArgbImage.py +++ b/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) - iType = IcnsType.match_maxsize(len(idat), 'argb' if is_argb else 'rgb') + uncompressed_data = PackBytes.unpack(data) - self.size = iType.size - self.channels = iType.channels or 0 - self.a, self.r, self.g, self.b = iType.split_channels(idat) + self.channels = 4 if is_argb else 3 + per_channel = len(uncompressed_data) // self.channels + 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) + \ - b''.join(self._raw_rgb_channels(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)) + return b'ARGB' + self.mask_data(compress=compress) \ + + self.rgb_data(compress=compress) def _load_png(self, fname: str) -> None: if not PIL_ENABLED: diff --git a/icnsutil/IcnsType.py b/icnsutil/IcnsType.py index 9ade5bb..f011278 100644 --- a/icnsutil/IcnsType.py +++ b/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) - 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 + # Filter attributes desc = None - if ext == 'icns' and filename: - for candidate in ['template', 'selected', 'dark']: - if filename.endswith(candidate + '.icns'): - desc = candidate - break + size = None + maxsize = None + retina = False + + # 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 - raise CanNotDetermine( - 'Could not determine type for file: "{}" – one of {}.'.format( - filename, [x.key for x in choices])) + choices = best_choice + + raise CanNotDetermine('Could not determine type – one of {}.'.format( + [x.key for x in choices]))