Moved some classes to different files, fixed Atari dos tests

* consolidated bytes_per_sector and sector_size (which mean the same thing) into sector_size
* moved AtrHeader, XfdHeader to ataridos.py
* moved base classes like WriteableSector, SectorList, etc. to utils.py

atrcopy/__init__.py

@@ -8,9 +8,9 @@ except ImportError:
     raise RuntimeError("atrcopy %s requires numpy" % __version__)
 
 from errors import *
-from ataridos import AtariDosDiskImage, AtariDosFile, get_xex
+from ataridos import AtrHeader, AtariDosDiskImage, AtariDosFile, get_xex, add_atr_header
 from dos33 import Dos33DiskImage
-from diskimages import AtrHeader, BootDiskImage, add_atr_header
+from diskimages import BootDiskImage
 from kboot import KBootImage, add_xexboot_header
 from segments import SegmentData, SegmentSaver, DefaultSegment, EmptySegment, ObjSegment, RawSectorsSegment, user_bit_mask, match_bit_mask, comment_bit_mask, data_style, selected_bit_mask, diff_bit_mask, not_user_bit_mask, interleave_segments
 from spartados import SpartaDosDiskImage

atrcopy/ataridos.py

@@ -1,9 +1,9 @@
 import numpy as np
 
 from errors import *
-from diskimages import DiskImageBase, Directory, VTOC, WriteableSector, BaseSectorList
+from diskimages import DiskImageBase, BaseHeader
 from segments import EmptySegment, ObjSegment, RawSectorsSegment, DefaultSegment, SegmentSaver
-from utils import to_numpy
+from utils import *
 
 import logging
 log = logging.getLogger(__name__)
@@ -37,7 +37,7 @@ class AtariDosVTOC(VTOC):
         log.debug("vtoc after: %s" % self.sector_map[0:720])
         packed = np.packbits(self.sector_map[0:720])
         self.vtoc1[0x0a:0x64] = packed
-        s = WriteableSector(self.bytes_per_sector, self.vtoc1)
+        s = WriteableSector(self.sector_size, self.vtoc1)
         s.sector_num = 360
         self.sectors.append(s)
 
@@ -165,10 +165,10 @@ class AtariDosDirent(object):
         return True
     
     def get_sector_list(self, image):
-        sector_list = BaseSectorList(image.bytes_per_sector)
+        sector_list = BaseSectorList(image.header.sector_size)
         self.start_read(image)
         while True:
-            sector = WriteableSector(image.bytes_per_sector, None, self.current_sector)
+            sector = WriteableSector(image.header.sector_size, None, self.current_sector)
             sector_list.append(sector)
             _, last, _, _ = self.read_sector(image)
             if last:
@@ -213,6 +213,7 @@ class AtariDosDirent(object):
         self.file_num = index
         self.dos_2 = True
         self.in_use = True
+        log.debug("set_values: %s" % self)
 
 
 class MydosDirent(AtariDosDirent):
@@ -290,6 +291,118 @@ class AtariDosFile(object):
             pos += 4 + count
 
 
+class AtrHeader(BaseHeader):
+    sector_class = AtariDosWriteableSector
+
+    # ATR Format described in http://www.atarimax.com/jindroush.atari.org/afmtatr.html
+    format = np.dtype([
+        ('wMagic', '<u2'),
+        ('wPars', '<u2'),
+        ('wSecSize', '<u2'),
+        ('btParsHigh', 'u1'),
+        ('dwCRC','<u4'),
+        ('unused','<u4'),
+        ('btFlags','u1'),
+        ])
+    file_format = "ATR"
+    
+    def __init__(self, bytes=None, sector_size=128, initial_sectors=3, create=False):
+        BaseHeader.__init__(self, sector_size, initial_sectors, 360)
+        if create:
+            self.header_offset = 16
+            self.check_size(0)
+        if bytes is None:
+            return
+        
+        if len(bytes) == 16:
+            values = bytes.view(dtype=self.format)[0]
+            if values[0] != 0x296:
+                raise InvalidAtrHeader
+            self.image_size = (int(values[3]) * 256 * 256 + int(values[1])) * 16
+            self.sector_size = int(values[2])
+            self.crc = int(values[4])
+            self.unused = int(values[5])
+            self.flags = int(values[6])
+            self.header_offset = 16
+        else:
+            raise InvalidAtrHeader
+    
+    def __str__(self):
+        return "%s Disk Image (size=%d (%dx%db), crc=%d flags=%d unused=%d)" % (self.file_format, self.image_size, self.max_sectors, self.sector_size, self.crc, self.flags, self.unused)
+    
+    def encode(self, raw):
+        values = raw.view(dtype=self.format)[0]
+        values[0] = 0x296
+        paragraphs = self.image_size / 16
+        parshigh, pars = divmod(paragraphs, 256*256)
+        values[1] = pars
+        values[2] = self.sector_size
+        values[3] = parshigh
+        values[4] = self.crc
+        values[5] = self.unused
+        values[6] = self.flags
+        return raw
+
+    def check_size(self, size):
+        if size == 92160 or size == 92176:
+            self.image_size = 92160
+            self.sector_size = 128
+            self.initial_sector_size = 0
+            self.num_initial_sectors = 0
+        elif size == 184320 or size == 184336:
+            self.image_size = 184320
+            self.sector_size = 256
+            self.initial_sector_size = 0
+            self.num_initial_sectors = 0
+        elif size == 183936 or size == 183952:
+            self.image_size = 183936
+            self.sector_size = 256
+            self.initial_sector_size = 128
+            self.num_initial_sectors = 3
+        else:
+            self.image_size = size
+        self.first_vtoc = 360
+        self.num_vtoc = 1
+        self.first_directory = 361
+        self.num_directory = 8
+        self.tracks_per_disk = 40
+        self.sectors_per_track = 18
+        initial_bytes = self.initial_sector_size * self.num_initial_sectors
+        self.max_sectors = ((self.image_size - initial_bytes) / self.sector_size) + self.num_initial_sectors
+    
+    def get_pos(self, sector):
+        if not self.sector_is_valid(sector):
+            raise ByteNotInFile166("Sector %d out of range" % sector)
+        if sector <= self.num_initial_sectors:
+            pos = self.num_initial_sectors * (sector - 1)
+            size = self.initial_sector_size
+        else:
+            pos = self.num_initial_sectors * self.initial_sector_size + (sector - 1 - self.num_initial_sectors) * self.sector_size
+            size = self.sector_size
+        pos += self.header_offset
+        return pos, size
+
+
+class XfdHeader(AtrHeader):
+    file_format = "XFD"
+    
+    def __str__(self):
+        return "%s Disk Image (size=%d (%dx%db)" % (self.file_format, self.image_size, self.max_sectors, self.sector_size)
+    
+    def __len__(self):
+        return 0
+    
+    def to_array(self):
+        raw = np.zeros([0], dtype=np.uint8)
+        return raw
+
+    def strict_check(self, image):
+        size = len(image)
+        if size in [92160, 133120, 183936, 184320]:
+            return
+        raise InvalidDiskImage("Uncommon size of XFD file")
+
+
 class AtariDosDiskImage(DiskImageBase):
     def __init__(self, *args, **kwargs):
         self.first_vtoc = 360
@@ -321,11 +434,37 @@ class AtariDosDiskImage(DiskImageBase):
     def __str__(self):
         return "%s Atari DOS Format: %d usable sectors (%d free), %d files" % (self.header, self.total_sectors, self.unused_sectors, len(self.files))
     
+    @classmethod
+    def new_header(cls, diskimage, format="ATR"):
+        if format.lower() == "atr":
+            header = AtrHeader(create=True)
+            header.check_size(diskimage.size)
+        else:
+            raise RuntimeError("Unknown header type %s" % format)
+        return header
+    
+    def as_new_format(self, format="ATR"):
+        """ Create a new disk image in the specified format
+        """
+        first_data = len(self.header)
+        raw = self.rawdata[first_data:]
+        data = add_atr_header(raw)
+        newraw = SegmentData(data)
+        image = self.__class__(newraw)
+        return image
+    
     vtoc_type = np.dtype([
         ('code', 'u1'),
         ('total','<u2'),
         ('unused','<u2'),
         ])
+    
+    def read_header(self):
+        bytes = self.bytes[0:16]
+        try:
+            self.header = AtrHeader(bytes)
+        except InvalidAtrHeader:
+            self.header = XfdHeader()
 
     def calc_vtoc_code(self):
         # From AA post: http://atariage.com/forums/topic/179868-mydos-vtoc-size/
@@ -495,3 +634,12 @@ def get_xex(segments, runaddr):
     words[1] = 0x2e1
     words[2] = runaddr
     return bytes
+
+def add_atr_header(bytes):
+    header = AtrHeader(create=True)
+    header.check_size(len(bytes))
+    hlen = len(header)
+    data = np.empty([hlen + len(bytes)], dtype=np.uint8)
+    data[0:hlen] = header.to_array()
+    data[hlen:] = bytes
+    return data

atrcopy/diskimages.py

@@ -77,117 +77,6 @@ class BaseHeader(object):
         return self.sector_class(self.sector_size, data)
 
 
-
-class AtrHeader(BaseHeader):
-    # ATR Format described in http://www.atarimax.com/jindroush.atari.org/afmtatr.html
-    format = np.dtype([
-        ('wMagic', '<u2'),
-        ('wPars', '<u2'),
-        ('wSecSize', '<u2'),
-        ('btParsHigh', 'u1'),
-        ('dwCRC','<u4'),
-        ('unused','<u4'),
-        ('btFlags','u1'),
-        ])
-    file_format = "ATR"
-    
-    def __init__(self, bytes=None, sector_size=128, initial_sectors=3, create=False):
-        BaseHeader.__init__(self, sector_size, initial_sectors, 360)
-        if create:
-            self.header_offset = 16
-            self.check_size(0)
-        if bytes is None:
-            return
-        
-        if len(bytes) == 16:
-            values = bytes.view(dtype=self.format)[0]
-            if values[0] != 0x296:
-                raise InvalidAtrHeader
-            self.image_size = (int(values[3]) * 256 * 256 + int(values[1])) * 16
-            self.sector_size = int(values[2])
-            self.crc = int(values[4])
-            self.unused = int(values[5])
-            self.flags = int(values[6])
-            self.header_offset = 16
-        else:
-            raise InvalidAtrHeader
-    
-    def __str__(self):
-        return "%s Disk Image (size=%d (%dx%db), crc=%d flags=%d unused=%d)" % (self.file_format, self.image_size, self.max_sectors, self.sector_size, self.crc, self.flags, self.unused)
-    
-    def encode(self, raw):
-        values = raw.view(dtype=self.format)[0]
-        values[0] = 0x296
-        paragraphs = self.image_size / 16
-        parshigh, pars = divmod(paragraphs, 256*256)
-        values[1] = pars
-        values[2] = self.sector_size
-        values[3] = parshigh
-        values[4] = self.crc
-        values[5] = self.unused
-        values[6] = self.flags
-        return raw
-
-    def check_size(self, size):
-        if size == 92160 or size == 92176:
-            self.image_size = 92160
-            self.sector_size = 128
-            self.initial_sector_size = 0
-            self.num_initial_sectors = 0
-        elif size == 184320 or size == 184336:
-            self.image_size = 184320
-            self.sector_size = 256
-            self.initial_sector_size = 0
-            self.num_initial_sectors = 0
-        elif size == 183936 or size == 183952:
-            self.image_size = 183936
-            self.sector_size = 256
-            self.initial_sector_size = 128
-            self.num_initial_sectors = 3
-        else:
-            self.image_size = size
-        self.first_vtoc = 360
-        self.num_vtoc = 1
-        self.first_directory = 361
-        self.num_directory = 8
-        self.tracks_per_disk = 40
-        self.sectors_per_track = 18
-        initial_bytes = self.initial_sector_size * self.num_initial_sectors
-        self.max_sectors = ((self.image_size - initial_bytes) / self.sector_size) + self.num_initial_sectors
-    
-    def get_pos(self, sector):
-        if not self.sector_is_valid(sector):
-            raise ByteNotInFile166("Sector %d out of range" % sector)
-        if sector <= self.num_initial_sectors:
-            pos = self.num_initial_sectors * (sector - 1)
-            size = self.initial_sector_size
-        else:
-            pos = self.num_initial_sectors * self.initial_sector_size + (sector - 1 - self.num_initial_sectors) * self.sector_size
-            size = self.sector_size
-        pos += self.header_offset
-        return pos, size
-
-
-class XfdHeader(AtrHeader):
-    file_format = "XFD"
-    
-    def __str__(self):
-        return "%s Disk Image (size=%d (%dx%db)" % (self.file_format, self.image_size, self.max_sectors, self.sector_size)
-    
-    def __len__(self):
-        return 0
-    
-    def to_array(self):
-        raw = np.zeros([0], dtype=np.uint8)
-        return raw
-
-    def strict_check(self, image):
-        size = len(image)
-        if size in [92160, 133120, 183936, 184320]:
-            return
-        raise InvalidDiskImage("Uncommon size of XFD file")
-
-
 class DiskImageBase(object):
     def __init__(self, rawdata, filename=""):
         self.rawdata = rawdata
@@ -273,22 +162,12 @@ class DiskImageBase(object):
     
     @classmethod
     def new_header(cls, diskimage, format="ATR"):
-        if format.lower() == "atr":
+        raise NotImplementedError
-            header = AtrHeader(create=True)
-            header.check_size(diskimage.size)
-        else:
-            raise RuntimeError("Unknown header type %s" % format)
-        return header
     
     def as_new_format(self, format="ATR"):
         """ Create a new disk image in the specified format
         """
-        first_data = len(self.header)
+        raise NotImplementedError
-        raw = self.rawdata[first_data:]
-        data = add_atr_header(raw)
-        newraw = SegmentData(data)
-        image = self.__class__(newraw)
-        return image
     
     def save(self, filename=""):
         if not filename:
@@ -310,11 +189,7 @@ class DiskImageBase(object):
             raise InvalidDiskImage("Invalid directory entries; may be boot disk")
     
     def read_header(self):
-        bytes = self.bytes[0:16]
+        raise NotImplementedError
-        try:
-            self.header = AtrHeader(bytes)
-        except InvalidAtrHeader:
-            self.header = XfdHeader()
     
     def check_size(self):
         pass
@@ -461,7 +336,7 @@ class DiskImageBase(object):
             data = to_numpy(data)
             sector_list = self.sector_builder_class(self.header, self.payload_bytes_per_sector, data, self.writeable_sector_class)
             vtoc_segments = self.get_vtoc_segments()
-            vtoc = self.vtoc_class(self.bytes_per_sector, vtoc_segments)
+            vtoc = self.vtoc_class(self.header, vtoc_segments)
             directory.save_dirent(self, dirent, vtoc, sector_list)
             self.write_sector_list(sector_list)
             self.write_sector_list(vtoc)
@@ -476,7 +351,7 @@ class DiskImageBase(object):
     def write_sector_list(self, sector_list):
         for sector in sector_list:
             pos, size = self.header.get_pos(sector.sector_num)
-            log.debug("writing: %s" % sector)
+            log.debug("writing: %s at %d" % (sector, pos))
             self.bytes[pos:pos + size] = sector.data
 
     def delete_file(self, filename):
@@ -527,12 +402,3 @@ class BootDiskImage(DiskImageBase):
             raise InvalidDiskImage("Number of boot sectors out of range")
         if bload < 0x200 or bload > (0xc000 - (nsec * self.header.sector_size)):
             raise InvalidDiskImage("Bad boot load address")
-
-def add_atr_header(bytes):
-    header = AtrHeader(create=True)
-    header.check_size(len(bytes))
-    hlen = len(header)
-    data = np.empty([hlen + len(bytes)], dtype=np.uint8)
-    data[0:hlen] = header.to_array()
-    data[hlen:] = bytes
-    return data

atrcopy/dos33.py

@@ -68,7 +68,7 @@ class Dos33VTOC(VTOC):
 
         # FIXME
         self.vtoc[0x38:] = vtoc.flatten()
-        s = WriteableSector(self.bytes_per_sector, self.vtoc)
+        s = WriteableSector(self.sector_size, self.vtoc)
         s.sector_num = 17 * 16
         self.sectors.append(s)
 
@@ -79,7 +79,7 @@ class Dos33Directory(Directory):
         return Dos33Dirent
 
     def get_dirent_sector(self):
-        s = self.sector_class(self.bytes_per_sector)
+        s = self.sector_class(self.sector_size)
         data = np.zeros([0x0b], dtype=np.uint8)
         s.add_data(data)
         return s
@@ -232,16 +232,16 @@ class Dos33Dirent(object):
         self.sector_map = sector_list
     
     def get_sector_list(self, image):
-        sector_list = BaseSectorList(image.header.bytes_per_sector)
+        sector_list = BaseSectorList(image.header.sector_size)
         self.start_read(image)
         sector_num = image.header.sector_from_track(self.track, self.sector)
         while sector_num > 0:
-            sector = WriteableSector(image.header.bytes_per_sector, None, sector_num)
+            sector = WriteableSector(image.header.sector_size, None, sector_num)
             sector_list.append(sector)
             values, style = image.get_sectors(sector_num)
             sector = image.header.sector_from_track(values[1], values[2])
         for sector_num in sector_list:
-            sector = WriteableSector(image.header.bytes_per_sector, None, sector_num)
+            sector = WriteableSector(image.header.sector_size, None, sector_num)
             sector_list.append(sector)
         return sector_list
 
@@ -320,11 +320,11 @@ class Dos33DiskImage(DiskImageBase):
         self.header = Dos33Header()
 
     @property
-    def bytes_per_sector(self):
+    def sector_size(self):
         return 256
 
     @property
-    def payload_bytes_per_sector(self):
+    def payload_sector_size(self):
         return 256
 
     @property
@@ -376,14 +376,14 @@ class Dos33DiskImage(DiskImageBase):
         ('unused5', 'S2'),
         ('num_tracks', 'u1'),
         ('sectors_per_track', 'u1'),
-        ('bytes_per_sector', 'u2'),
+        ('sector_size', 'u2'),
         ])
 
     def get_vtoc(self):
         data, style = self.get_sectors(self.header.first_vtoc)
         values = data[0:self.vtoc_type.itemsize].view(dtype=self.vtoc_type)[0]
         self.header.first_directory = self.header.sector_from_track(values['cat_track'], values['cat_sector'])
-        self.header.sector_size = int(values['bytes_per_sector'])
+        self.header.sector_size = int(values['sector_size'])
         self.header.max_sectors = int(values['num_tracks']) * int(values['sectors_per_track'])
         self.header.ts_pairs = int(values['max_pairs'])
         self.header.dos_release = values['dos_release']

atrcopy/utils.py

@@ -2,6 +2,9 @@ import types
 
 import numpy as np
 
+import logging
+log = logging.getLogger(__name__)
+
 
 def to_numpy(value):
     if type(value) is np.ndarray:
@@ -69,8 +72,8 @@ class WriteableSector(object):
 
 
 class BaseSectorList(object):
-    def __init__(self, bytes_per_sector):
+    def __init__(self, sector_size):
-        self.bytes_per_sector = bytes_per_sector
+        self.sector_size = sector_size
         self.sectors = []
 
     def __len__(self):
@@ -97,7 +100,7 @@ class BaseSectorList(object):
 
 class Directory(BaseSectorList):
     def __init__(self, header, num_dirents=-1, sector_class=WriteableSector):
-        BaseSectorList.__init__(self, header.bytes_per_sector)
+        BaseSectorList.__init__(self, header.sector_size)
         self.sector_class = sector_class
         self.num_dirents = num_dirents
         # number of dirents may be unlimited, so use a dict instead of a list
@@ -193,7 +196,7 @@ class Directory(BaseSectorList):
         self.finish_encoding(image)
 
     def get_dirent_sector(self):
-        return self.sector_class(self.bytes_per_sector)
+        return self.sector_class(self.sector_size)
 
     def encode_empty(self):
         raise NotImplementedError
@@ -222,7 +225,7 @@ class Directory(BaseSectorList):
 
 class VTOC(BaseSectorList):
     def __init__(self, header, segments=None):
-        BaseSectorList.__init__(self, header.bytes_per_sector)
+        BaseSectorList.__init__(self, header.sector_size)
 
         # sector map: 1 is free, 0 is allocated
         self.sector_map = np.zeros([1280], dtype=np.uint8)
@@ -259,7 +262,7 @@ class VTOC(BaseSectorList):
 
 class SectorBuilder(BaseSectorList):
     def __init__(self, header, usable, data, sector_class):
-        BaseSectorList.__init__(self, header.bytes_per_sector)
+        BaseSectorList.__init__(self, header.sector_size)
         self.data = to_numpy(data)
         self.usable_bytes = usable
         self.split_into_sectors(header)

test/test_add_file.py

@@ -121,7 +121,7 @@ class BaseFilesystemModifyTest(object):
 # class TestAtariDosSDImage(BaseFilesystemModifyTest):
 #     diskimage_type = AtariDosDiskImage
 #     sample_file = "../test_data/dos_sd_test1.atr"
-#     num_files_in_sample
+#     num_files_in_sample = 5
 
 class TestDos33Image(BaseFilesystemModifyTest):
     diskimage_type = Dos33DiskImage