fix adding nested icns file (incl. test)

icnsutil/IcnsFile.py

@@ -8,6 +8,8 @@ from .ArgbImage import ArgbImage
 
 
 class IcnsFile:
+    __slots__ = ['media', 'infile']
+
     @staticmethod
     def verify(fname: str) -> Iterator[str]:
         '''
@@ -129,6 +131,9 @@ class IcnsFile:
                 print('Warning: unknown media type: {}, {} bytes, "{}"'.format(
                     str(key), len(data), file), file=stderr)
 
+    def has_toc(self) -> bool:
+        return 'TOC ' in self.media.keys()
+
     def add_media(self, key: Optional[IcnsType.Media.KeyT] = None, *,
                   file: Optional[str] = None, data: Optional[bytes] = None,
                   force: bool = False) -> None:
@@ -143,14 +148,29 @@ class IcnsFile:
                 data = fp.read()
         if not data:
             raise AttributeError('Did you miss file= or data= attribute?')
+
         if not key:  # Determine ICNS type
-            key = IcnsType.guess(data, file).key
+            iType = IcnsType.guess(data, file)
+            key = iType.key
+            is_icns = iType.is_type('icns')
+        else:
+            is_icns = True  # we dont know, so we assume it is
+
         # Check if type is unique
         if not force and key in self.media.keys():
             raise KeyError('Image with identical key "{}". File: {}'.format(
                 str(key), file))
+        # Nested icns files must omit the icns header
+        if is_icns and data[:4] == b'icns':
+            data = data[8:]
         self.media[key] = data
 
+    def remove_media(self, key: IcnsType.Media.KeyT) -> bool:
+        if key not in self.media.keys():
+            return False
+        del self.media[key]
+        return True
+
     def write(self, fname: str, *, toc: bool = True) -> None:
         ''' Create a new ICNS file from stored media. '''
         # Rebuild TOC to ensure soundness
@@ -232,8 +252,8 @@ class IcnsFile:
 
     def _make_toc(self, *, enabled: bool) -> List[IcnsType.Media.KeyT]:
         # Rebuild TOC to ensure soundness
-        if 'TOC ' in self.media.keys():
+        if self.has_toc():
-            del(self.media['TOC '])
+            del self.media['TOC ']
         # We loop two times over the keys; so, make sure order is identical.
         # By default this will be the same order as read/written.
         order = list(self.media.keys())

tests/test_icnsutil.py

@@ -171,10 +171,40 @@ class TestIcnsFile(unittest.TestCase):
         self.assertEqual(
             list(newimg.media.keys()), ['icp5', 'ic11', 'ic08', 'ic13'])
 
+    def test_add_nested_icns(self):
+        img = IcnsFile()
+        img.add_media(file='selected.icns')
+        self.assertTrue('slct' in img.media.keys())
+        self.assertNotEqual(img.media['slct'][:4], b'icns')
+        with self.assertRaises(KeyError):
+            img.add_media('slct', file='rgb.icns')
+        with self.assertRaises(IcnsType.CanNotDetermine):
+            img.add_media(file='icp4rgb.icns')
+        with self.assertRaises(IcnsType.CanNotDetermine):
+            img.add_media(file='rgb.icns')
+        img.add_media('no_key', file='rgb.icns')
+        self.assertTrue('no_key' in img.media.keys())
+        self.assertNotEqual(img.media['no_key'][:4], b'icns')
+
+    def test_remove_media(self):
+        img = IcnsFile()
+        img.add_media(file='selected.icns')
+        img.add_media(file='rgb.icns.rgb')
+        img.add_media(file='rgb.icns.argb')
+        self.assertListEqual(list(img.media.keys()), ['slct', 'is32', 'ic04'])
+        img.remove_media('is32')
+        self.assertListEqual(list(img.media.keys()), ['slct', 'ic04'])
+        img.remove_media(b'ic04')
+        self.assertListEqual(list(img.media.keys()), ['slct', 'ic04'])
+        img.remove_media('')
+        self.assertListEqual(list(img.media.keys()), ['slct', 'ic04'])
+        img.remove_media('slct')
+        self.assertListEqual(list(img.media.keys()), ['ic04'])
+
     def test_toc(self):
         img = IcnsFile()
         fname_out = 'tmp-out.icns'
-        img.add_media(file='rgb.icns.argb', key='ic04')
+        img.add_media('ic04', file='rgb.icns.argb')
         # without TOC
         img.write(fname_out, toc=False)
         with open(fname_out, 'rb') as fp:
@@ -183,6 +213,7 @@ class TestIcnsFile(unittest.TestCase):
             self.assertEqual(fp.read(4), b'ic04')
             self.assertEqual(fp.read(4), b'\x00\x00\x02\xD1')
             self.assertEqual(fp.read(4), b'ARGB')
+        self.assertFalse(IcnsFile(fname_out).has_toc())
         # with TOC
         img.write(fname_out, toc=True)
         with open(fname_out, 'rb') as fp:
@@ -195,6 +226,7 @@ class TestIcnsFile(unittest.TestCase):
             self.assertEqual(fp.read(4), b'ic04')
             self.assertEqual(fp.read(4), b'\x00\x00\x02\xD1')
             self.assertEqual(fp.read(4), b'ARGB')
+        self.assertTrue(IcnsFile(fname_out).has_toc())
         os.remove(fname_out)
 
     def test_verify(self):