Added bz2, xz containers, tests, sample data

README.rst

@@ -135,6 +135,10 @@ running ``atrcopy``.
 +====================+==========+======+=======+==============================+
 | gzip               | .gz      | Yes  | No    | Read only                    |
 +--------------------+----------+------+-------+------------------------------+
+| bzip2              | .bz2     | Yes  | No    | Read only                    |
++--------------------+----------+------+-------+------------------------------+
+| lzma               | .xz      | Yes  | No    | Read only                    |
++--------------------+----------+------+-------+------------------------------+
 | Disk Communicator  | .dcm     | No   | No    | Recognized but unimplemented |
 +--------------------+----------+------+-------+------------------------------+
 

atrcopy/__init__.py

@@ -20,7 +20,7 @@ from .kboot import KBootImage, add_xexboot_header
 from .segments import SegmentData, SegmentSaver, DefaultSegment, EmptySegment, ObjSegment, RawSectorsSegment, SegmentedFileSegment, user_bit_mask, match_bit_mask, comment_bit_mask, data_style, selected_bit_mask, diff_bit_mask, not_user_bit_mask, interleave_segments, SegmentList, get_style_mask, get_style_bits
 from .spartados import SpartaDosDiskImage
 from .cartridge import A8CartHeader, AtariCartImage
-from .parsers import SegmentParser, DefaultSegmentParser, guess_parser_for_mime, guess_parser_for_system, guess_container, iter_known_segment_parsers, mime_parse_order, parsers_for_filename
+from .parsers import SegmentParser, DefaultSegmentParser, guess_parser_for_mime, guess_parser_for_system, guess_container, iter_parsers, iter_known_segment_parsers, mime_parse_order, parsers_for_filename
 from .magic import guess_detail_for_mime
 from .utils import to_numpy, text_to_int
 

atrcopy/container.py

@@ -1,4 +1,6 @@
 import gzip
+import bz2
+import lzma
 import io
 
 import numpy as np
@@ -53,3 +55,25 @@ class GZipContainer(DiskImageContainer):
         except OSError as e:
             raise errors.InvalidContainer(e)
         return unpacked
+
+
+class BZipContainer(DiskImageContainer):
+    def unpack_bytes(self, byte_data):
+        try:
+            buf = io.BytesIO(byte_data)
+            with bz2.BZ2File(buf, mode='rb') as f:
+                unpacked = f.read()
+        except OSError as e:
+            raise errors.InvalidContainer(e)
+        return unpacked
+
+
+class LZMAContainer(DiskImageContainer):
+    def unpack_bytes(self, byte_data):
+        try:
+            buf = io.BytesIO(byte_data)
+            with lzma.LZMAFile(buf, mode='rb') as f:
+                unpacked = f.read()
+        except lzma.LZMAError as e:
+            raise errors.InvalidContainer(e)
+        return unpacked

atrcopy/parsers.py

@@ -10,7 +10,7 @@ from .dos33 import Dos33DiskImage, ProdosDiskImage, Dos33BinFile
 from .standard_delivery import StandardDeliveryImage
 from . import errors
 from .magic import guess_detail_for_mime
-from .container import GZipContainer
+from . import container
 from .dcm import DCMContainer
 
 import logging
@@ -160,7 +160,9 @@ class ProdosSegmentParser(SegmentParser):
 
 
 known_containers = [
-    GZipContainer,
+    container.GZipContainer,
+    container.BZipContainer,
+    container.LZMAContainer,
     DCMContainer,
 ]
 
@@ -207,7 +209,7 @@ def guess_parser_for_system(mime_base, r):
 def iter_parsers(r):
     container = guess_container(r.data)
     if container is not None:
-        r = container.unpacked
+        r = SegmentData(container.unpacked)
     for mime in mime_parse_order:
         p = guess_parser_for_mime(mime, r)
         if p is not None:

test/test_container.py

@@ -0,0 +1,41 @@
+from __future__ import print_function
+from builtins import object
+import numpy as np
+
+from mock import *
+
+from atrcopy import SegmentData, iter_parsers
+from atrcopy import errors
+
+
+class BaseContainerTest:
+    base_path = None
+    expected_mime = ""
+
+    @pytest.mark.parametrize("ext", ['.gz', '.bz2', '.xz', '.dcm'])
+    def test_container(self, ext):
+        pathname = self.base_path + ext
+        try:
+            sample_data = np.fromfile(pathname, dtype=np.uint8)
+        except OSError:
+            pass
+        else:
+            rawdata = SegmentData(sample_data.copy())
+            mime, parser = iter_parsers(rawdata)
+            assert mime == self.expected_mime
+            assert len(parser.image.files) == self.num_files_in_sample
+
+class TestContainerAtariDosSDImage(BaseContainerTest):
+    base_path = "../test_data/container_dos_sd_test1.atr"
+    expected_mime = "application/vnd.atari8bit.atr"
+    num_files_in_sample = 5
+
+class TestContainerAtariDosEDImage(BaseContainerTest):
+    base_path = "../test_data/container_dos_ed_test1.atr"
+    expected_mime = "application/vnd.atari8bit.atr"
+    num_files_in_sample = 5
+
+class TestContainerAtariDosDDImage(BaseContainerTest):
+    base_path = "../test_data/container_dos_dd_test1.atr"
+    expected_mime = "application/vnd.atari8bit.atr"
+    num_files_in_sample = 5