Added sector convenience functions that return segments

atrcopy/container.py

@@ -319,3 +319,10 @@ def guess_container(raw_data):
         log.debug(f"image does not appear to be compressed.")
         container = Container(raw_data)
     return container
+
+
+def load(pathname):
+    sample_data = np.fromfile(pathname, dtype=np.uint8)
+    container = guess_container(sample_data)
+    container.guess_media_type()
+    return container

atrcopy/media_type.py

@@ -125,7 +125,21 @@ class DiskImage(MediaType):
         if not self.sector_is_valid(sector):
             raise errors.ByteNotInFile166("Sector %d out of range" % sector)
         pos = (sector - self.starting_sector_label) * self.sector_size
-        return pos + self.header_length, self.sector_size
+        return pos, self.sector_size
+
+    def get_contiguous_sectors(self, start, count):
+        index, _ = self.get_index_of_sector(start)
+        last, size = self.get_index_of_sector(start + count - 1)
+        return Segment(self, index, length=(last + size - index))
+
+    def get_sector_list(self, sector_numbers):
+        offsets = np.empty(len(sector_numbers) * self.sector_size, dtype=np.uint32)
+        i = 0
+        for num in sector_numbers:
+            index, size = self.get_index_of_sector(num)
+            offsets[i:i+size] = np.arange(index, index + size)
+            i += size
+        return Segment(self, offsets)
 
 
 class CartImage(MediaType):

test/test_disk_image.py

@@ -0,0 +1,65 @@
+import glob
+
+import numpy as np
+
+from mock import *
+
+from atrcopy.container import guess_container, load
+from atrcopy.media_type import MediaType, guess_media_type
+from atrcopy import errors
+
+from atrcopy.media_types.atari_disks import *
+from atrcopy.media_types.apple_disks import *
+
+class TestDiskImage:
+
+    @pytest.mark.parametrize("filename,sector_size,sectors", [
+        ("dos_sd_test1.atr", 128, [(1,3), (360,1), (361,8)]),
+        ("dos_ed_test1.atr", 128, [(1,3), (360,1), (361,8)]),
+    ])
+    def test_contiguous_sectors(self, filename, sector_size, sectors):
+        pathname = os.path.join(os.path.dirname(__file__), "../test_data", filename)
+        container = load(pathname)
+        media = container.media
+
+        for start, count in sectors:
+            # grab sectors
+            segment = media.get_contiguous_sectors(start, count)
+            # 16 byte ATR header, first byte of sector 1 is byte 16
+            assert np.array_equal(segment.container_offset, np.arange(16 + (start - 1) * sector_size, 16 + (start + count - 1)*sector_size))
+
+    def test_ed_sectors(self):
+        pathname = os.path.join(os.path.dirname(__file__), "../test_data/dos_ed_test1.atr")
+        container = load(pathname)
+        media = container.media
+
+        # get vtoc + vtoc2
+        boot = media.get_sector_list([360,1024])
+        assert np.array_equal(boot.container_offset[0:128], np.arange(16 + (360 - 1) * 128, 16 + 360*128))
+        assert np.array_equal(boot.container_offset[128:256], np.arange(16 + (1024 - 1) * 128, 16 + 1024*128))
+
+
+if __name__ == "__main__":
+    import logging
+    logging.basicConfig(level=logging.WARNING)
+    log = logging.getLogger("atrcopy.media_type")
+    log.setLevel(logging.DEBUG)
+
+    def check(pathname):
+        print(f"checking {pathname}")
+        container = load(pathname)
+        print(container.verbose_info)
+        media = container.media
+        print(media)
+        boot = media.get_contiguous_sectors(1, 3)
+        print(boot)
+        print(boot.container_offset)
+
+    import sys
+    import glob
+    if len(sys.argv) > 1:
+        images = sys.argv[1:]
+    else:
+        images = sorted(glob.glob(os.path.join(os.path.dirname(__file__), "../test_data/", "*")))
+    for pathname in images:
+        check(pathname)