diff --git a/atrcopy/container.py b/atrcopy/container.py index 46bf798..3cf6818 100644 --- a/atrcopy/container.py +++ b/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 diff --git a/atrcopy/media_type.py b/atrcopy/media_type.py index fac6102..4590333 100644 --- a/atrcopy/media_type.py +++ b/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): diff --git a/test/test_disk_image.py b/test/test_disk_image.py new file mode 100644 index 0000000..d026119 --- /dev/null +++ b/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)