Fixed file number in Atari DOS files

atrcopy/ataridos.py

@@ -19,9 +19,9 @@ class AtariDosWriteableSector(WriteableSector):
         self._next_sector_num = value
         index = self.sector_size - 3
         hi, lo = divmod(value, 256)
-        self.data[index] = self.used
+        self.data[index + 0] = (self.file_num << 2) | (hi & 0x03)
         self.data[index + 1] = lo
-        self.data[index + 2] = hi
+        self.data[index + 2] = self.used
         log.debug("sector metadata for %d: %s" % (self._sector_num, self.data[index:index + 3]))
         # file number will be added later when known.
 
@@ -177,7 +177,7 @@ class AtariDosDirent(object):
     def process_raw_sector(self, image, raw):
         file_num = raw[-3] >> 2
         if file_num != self.file_num:
-            raise FileNumberMismatchError164()
+            raise FileNumberMismatchError164("Expecting file %d, found %d" % (self.file_num, file_num))
         self.sectors_seen.add(self.current_sector)
         next_sector = ((raw[-3] & 0x3) << 8) + raw[-2]
         if next_sector in self.sectors_seen:

atrcopy/diskimages.py

@@ -387,7 +387,7 @@ class DiskImageBase(object):
         sector_list = self.sector_list_class(self.bytes_per_sector, 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)
-        sector_list.calc_sector_map(vtoc)
+        sector_list.calc_sector_map(dirent, vtoc)
         directory.save_dirent(dirent, sector_list)
         self.write_sector_list(sector_list)
         self.write_sector_list(vtoc)
@@ -406,6 +406,7 @@ class WriteableSector(object):
         self._sector_num = -1
         self._next_sector = 0
         self.sector_size = sector_size
+        self.file_num = 0
         self.data = np.zeros([sector_size], dtype=np.uint8)
         self.used = 0
         self.ptr = self.used
@@ -612,7 +613,7 @@ class SectorList(BaseSectorList):
             self.sectors.append(sector)
             index += count
 
-    def calc_sector_map(self, vtoc):
+    def calc_sector_map(self, dirent, vtoc):
         """ Map out the sectors and link the sectors together
 
         raises NotEnoughSpaceOnDisk if the whole file won't fit. It will not
@@ -627,6 +628,7 @@ class SectorList(BaseSectorList):
         last_sector = None
         for sector, sector_num in zip(self.sectors, order):
             sector.sector_num = sector_num
+            sector.file_num = dirent.file_num
             self.file_length += sector.used
             if last_sector is not None:
                 last_sector.next_sector_num = sector_num

test/test_add_file.py

@@ -14,10 +14,13 @@ class TestAtariDosSDImage(object):
     def test_small(self):
         assert len(self.image.files) == 5
 
-        data = [0xff, 0xff, 0x00, 0x60, 0x01, 0x60, 1, 2]
+        data = np.asarray([0xff, 0xff, 0x00, 0x60, 0x01, 0x60, 1, 2], dtype=np.uint8)
         self.image.write_file("TEST.XEX", None, data)
         assert len(self.image.files) == 6
 
+        data2 = self.image.find_file("TEST.XEX")
+        assert data.tostring() == data2
+
 
 if __name__ == "__main__":
     t = TestAtariDosFile()