From c8ccd737796c873f65e035c4fafaae9293c8b877 Mon Sep 17 00:00:00 2001 From: Rob McMullen Date: Fri, 25 Mar 2016 16:07:16 -0700 Subject: [PATCH] Added SegmentData class and updated to version 2.5.0 in preparation for adding comments --- atrcopy/__init__.py | 18 ++++++------- atrcopy/ataridos.py | 50 +++++++++++++++++------------------ atrcopy/diskimages.py | 49 ++++++++++++++++++---------------- atrcopy/segments.py | 61 ++++++++++++++++++++++++++++++++++--------- atrcopy/spartados.py | 18 ++++++------- 5 files changed, 116 insertions(+), 80 deletions(-) diff --git a/atrcopy/__init__.py b/atrcopy/__init__.py index c506de7..1d3545c 100755 --- a/atrcopy/__init__.py +++ b/atrcopy/__init__.py @@ -1,4 +1,4 @@ -__version__ = "2.4.0" +__version__ = "2.5.0" try: import numpy as np @@ -9,7 +9,7 @@ from errors import * from ataridos import AtariDosDiskImage, AtariDosFile from diskimages import AtrHeader, BootDiskImage from kboot import KBootImage -from segments import SegmentSaver, DefaultSegment, EmptySegment, ObjSegment, RawSectorsSegment, IndexedByteSegment +from segments import SegmentData, SegmentSaver, DefaultSegment, EmptySegment, ObjSegment, RawSectorsSegment, IndexedByteSegment from spartados import SpartaDosDiskImage from utils import to_numpy @@ -60,12 +60,12 @@ def run(): for filename in options.files: with open(filename, "rb") as fh: - data = fh.read() + rawdata = SegmentData(fh.read()) + data = rawdata.get_data() image = None if options.debug: - data = to_numpy(data) header = AtrHeader(data[0:16]) - image = SpartaDosDiskImage(data, filename) + image = SpartaDosDiskImage(rawdata, filename) else: try: data = to_numpy(data) @@ -74,7 +74,7 @@ def run(): for format in [KBootImage, SpartaDosDiskImage, AtariDosDiskImage]: if options.verbose: print "trying", format.__name__ try: - image = format(data, filename) + image = format(rawdata, filename) print "%s: %s" % (filename, image) break except InvalidDiskImage: @@ -82,7 +82,7 @@ def run(): except AtrError: for format in [AtariDosDiskImage]: try: - image = format(data) + image = format(rawdata, filename) print "%s: %s" % (filename, image) break except: @@ -91,13 +91,13 @@ def run(): except AtrError: if options.verbose: print "%s: Doesn't look like a supported disk image" % filename try: - image = AtariDosFile(data) + image = AtariDosFile(rawdata) print "%s:\n%s" % (filename, image) except InvalidBinaryFile: if options.verbose: print "%s: Doesn't look like an XEX either" % filename continue if image is None: - image = BootDiskImage(data, filename) + image = BootDiskImage(rawdata, filename) if options.segments: image.parse_segments() print "\n".join([str(a) for a in image.segments]) diff --git a/atrcopy/ataridos.py b/atrcopy/ataridos.py index c254622..13ea71d 100755 --- a/atrcopy/ataridos.py +++ b/atrcopy/ataridos.py @@ -142,13 +142,11 @@ class AtariDosFile(object): Ref: http://www.atarimax.com/jindroush.atari.org/afmtexe.html """ - def __init__(self, data, style=None): - self.bytes = to_numpy(data) - self.size = np.alen(self.bytes) - if style is None: - self.style = np.zeros(self.size, dtype=np.uint8) - else: - self.style = style + def __init__(self, rawdata): + self.rawdata = rawdata + self.bytes = rawdata.get_data() + self.style = rawdata.get_style() + self.size = len(rawdata) self.segments = [] self.parse_segments() @@ -156,6 +154,7 @@ class AtariDosFile(object): return "\n".join(str(s) for s in self.segments) + "\n" def parse_segments(self): + r = self.rawdata b = self.bytes s = self.style pos = 0 @@ -164,7 +163,7 @@ class AtariDosFile(object): if pos + 1 < self.size: header, = b[pos:pos+2].view(dtype=' 0: start, count = self.get_contiguous_sectors(self.vtoc2, 1) - segment = RawSectorsSegment(b[start:start+count], s[start:start+count], self.vtoc2, 1, count, self.header.sector_size, name="VTOC2") + segment = RawSectorsSegment(r[start:start+count], self.vtoc2, 1, count, self.header.sector_size, name="VTOC2") segments.append(segment) return segments def get_directory_segments(self): - b = self.bytes - s = self.style + r = self.rawdata segments = [] addr = 0 start, count = self.get_contiguous_sectors(361, 8) - segment = RawSectorsSegment(b[start:start+count], s[start:start+count], 361, 8, count, 128, 3, self.header.sector_size, name="Directory") + segment = RawSectorsSegment(r[start:start+count], 361, 8, count, 128, 3, self.header.sector_size, name="Directory") segments.append(segment) return segments @@ -309,7 +307,7 @@ class AtariDosDiskImage(DiskImageBase): name = "%s %ds@%d" % (dirent.get_filename(), dirent.num_sectors, dirent.starting_sector) verbose_name = "%s (%d sectors, first@%d) %s" % (dirent.get_filename(), dirent.num_sectors, dirent.starting_sector, dirent.verbose_info) print verbose_name - segment = IndexedByteSegment(self.bytes, self.style, byte_order, name=name, verbose_name=verbose_name) + segment = IndexedByteSegment(self.rawdata, byte_order, name=name, verbose_name=verbose_name) else: - segment = EmptySegment(self.bytes, self.style, name=dirent.get_filename()) + segment = EmptySegment(self.rawdata, name=dirent.get_filename()) return segment diff --git a/atrcopy/diskimages.py b/atrcopy/diskimages.py index dab79b7..5173214 100755 --- a/atrcopy/diskimages.py +++ b/atrcopy/diskimages.py @@ -1,7 +1,7 @@ import numpy as np from errors import * -from segments import EmptySegment, ObjSegment, RawSectorsSegment, IndexedByteSegment +from segments import SegmentData, EmptySegment, ObjSegment, RawSectorsSegment from utils import to_numpy class AtrHeader(object): @@ -92,18 +92,11 @@ class XfdHeader(AtrHeader): class DiskImageBase(object): - debug = False - - def __init__(self, bytes, style=None, filename=""): - self.bytes = to_numpy(bytes) + def __init__(self, rawdata, filename=""): + self.rawdata = rawdata + self.bytes = self.rawdata.get_data() + self.style = self.rawdata.get_style() self.size = np.alen(self.bytes) - if style is None: - if self.debug: - self.style = np.arange(self.size, dtype=np.uint8) - else: - self.style = np.zeros(self.size, dtype=np.uint8) - else: - self.style = style self.set_filename(filename) self.header = None self.total_sectors = 0 @@ -164,7 +157,7 @@ class DiskImageBase(object): pos, size = self.header.get_pos(sector) return self.bytes[pos:pos + size], pos, size - def get_sectors(self, start, end=None): + def get_sector_slice(self, start, end=None): """ Get contiguous sectors :param start: first sector number to read (note: numbering starts from 1) @@ -178,7 +171,17 @@ class DiskImageBase(object): start += 1 _, more = self.header.get_pos(start) size += more - return self.bytes[pos:pos + size], self.style[pos:pos + size] + return slice(pos, pos + size) + + def get_sectors(self, start, end=None): + """ Get contiguous sectors + + :param start: first sector number to read (note: numbering starts from 1) + :param end: last sector number to read + :returns: bytes + """ + s = self.get_sector_slice(start, end) + return self.bytes[s], self.style[s] def get_contiguous_sectors(self, sector, num): start = 0 @@ -191,12 +194,11 @@ class DiskImageBase(object): return start, count def parse_segments(self): - b = self.bytes - s = self.style + r = self.rawdata i = self.header.atr_header_offset if self.header.image_size > 0: - self.segments.append(ObjSegment(b[0:i], s[0:i], 0, 0, 0, i, name="%s Header" % self.header.file_format)) - self.segments.append(RawSectorsSegment(b[i:], s[i:], 1, self.header.max_sectors, self.header.image_size, 128, 3, self.header.sector_size, name="Raw disk sectors")) + self.segments.append(ObjSegment(r[0:i], 0, 0, 0, i, name="%s Header" % self.header.file_format)) + self.segments.append(RawSectorsSegment(r[i:], 1, self.header.max_sectors, self.header.image_size, 128, 3, self.header.sector_size, name="Raw disk sectors")) self.segments.extend(self.get_boot_segments()) self.segments.extend(self.get_vtoc_segments()) self.segments.extend(self.get_directory_segments()) @@ -217,10 +219,11 @@ class DiskImageBase(object): if flag == 0: num = int(values[1]) addr = int(values[2]) - bytes, style = self.get_sectors(1, num) - header = ObjSegment(bytes[0:6], style[0:6], 0, 0, addr, addr + 6, name="Boot Header") - sectors = ObjSegment(bytes, style, 0, 0, addr, addr + len(bytes), name="Boot Sectors") - code = ObjSegment(bytes[6:], style[6:], 0, 0, addr + 6, addr + len(bytes), name="Boot Code") + s = self.get_sector_slice(1, num) + r = self.rawdata[s] + header = ObjSegment(r[0:6], 0, 0, addr, addr + 6, name="Boot Header") + sectors = ObjSegment(r, 0, 0, addr, addr + len(r), name="Boot Sectors") + code = ObjSegment(r[6:], 0, 0, addr + 6, addr + len(r), name="Boot Code") segments = [sectors, header, code] return segments @@ -249,7 +252,7 @@ class DiskImageBase(object): try: segment = self.get_file_segment(dirent) except InvalidFile, e: - segment = EmptySegment(self.bytes, self.style, name=dirent.get_filename(), error=str(e)) + segment = EmptySegment(self.rawdata, name=dirent.get_filename(), error=str(e)) segments.append(segment) return segments diff --git a/atrcopy/segments.py b/atrcopy/segments.py index 5f30cee..1f1426f 100755 --- a/atrcopy/segments.py +++ b/atrcopy/segments.py @@ -1,7 +1,7 @@ import numpy as np from errors import * -from utils import to_numpy_list +from utils import to_numpy, to_numpy_list class SegmentSaver(object): @@ -22,13 +22,47 @@ class SegmentSaver(object): return "|".join(wildcards) +class SegmentData(object): + def __init__(self, data, style=None, comments=None, debug=False): + self.data = to_numpy(data) + if style is None: + if debug: + self.style = np.arange(len(self), dtype=np.uint8) + else: + self.style = np.zeros(len(self), dtype=np.uint8) + else: + self.style = style + if comments is None: + comments = dict() + self.comments = comments + + def __len__(self): + return np.alen(self.data) + + def get_data(self): + return self.data + + def get_style(self): + return self.style + + def get_comments(self): + return self.comments + + def __getitem__(self, index): + d = self.data[index] + s = self.style[index] + c = self.comments + return SegmentData(d, s, c) + + class DefaultSegment(object): savers = [SegmentSaver] - def __init__(self, data, style, start_addr=0, name="All", error=None, verbose_name=None): + def __init__(self, rawdata, start_addr=0, name="All", error=None, verbose_name=None): self.start_addr = int(start_addr) # force python int to decouple from possibly being a numpy datatype - self.data = data - self.style = style + self.rawdata = rawdata + self.data = rawdata.get_data() + self.style = rawdata.get_style() self.error = error self.name = name self.verbose_name = verbose_name @@ -51,7 +85,7 @@ class DefaultSegment(object): return s def __len__(self): - return np.alen(self.data) + return len(self.rawdata) def __getitem__(self, index): return self.data[index] @@ -213,8 +247,8 @@ class DefaultSegment(object): class EmptySegment(DefaultSegment): - def __init__(self, data, style, name="", error=None): - DefaultSegment.__init__(self, data, style, 0, name, error) + def __init__(self, rawdata, name="", error=None): + DefaultSegment.__init__(self, rawdata, 0, name, error) def __str__(self): s = "%s (empty file)" % (self.name, ) @@ -234,8 +268,8 @@ class EmptySegment(DefaultSegment): class ObjSegment(DefaultSegment): - def __init__(self, data, style, metadata_start, data_start, start_addr, end_addr, name="", **kwargs): - DefaultSegment.__init__(self, data, style, start_addr, name, **kwargs) + def __init__(self, rawdata, metadata_start, data_start, start_addr, end_addr, name="", **kwargs): + DefaultSegment.__init__(self, rawdata, start_addr, name, **kwargs) self.metadata_start = metadata_start self.data_start = data_start @@ -257,8 +291,8 @@ class ObjSegment(DefaultSegment): class RawSectorsSegment(DefaultSegment): - def __init__(self, data, style, first_sector, num_sectors, count, boot_sector_size, num_boot_sectors, sector_size, **kwargs): - DefaultSegment.__init__(self, data, style, 0, **kwargs) + def __init__(self, rawdata, first_sector, num_sectors, count, boot_sector_size, num_boot_sectors, sector_size, **kwargs): + DefaultSegment.__init__(self, rawdata, 0, **kwargs) self.boot_sector_size = boot_sector_size self.num_boot_sectors = num_boot_sectors self.page_size = sector_size @@ -324,10 +358,11 @@ class IndexedStyleWrapper(object): class IndexedByteSegment(DefaultSegment): - def __init__(self, data, style, byte_order, **kwargs): + def __init__(self, rawdata, byte_order, **kwargs): # Convert to numpy list so fancy indexing works as argument to __getitem__ self.order = to_numpy_list(byte_order) - DefaultSegment.__init__(self, data, IndexedStyleWrapper(style, byte_order), 0, **kwargs) + DefaultSegment.__init__(self, rawdata, **kwargs) + self.style = IndexedStyleWrapper(self.style, byte_order) def __str__(self): s = "%s ($%x @ $%x)" % (self.name, len(self), self.order[0]) diff --git a/atrcopy/spartados.py b/atrcopy/spartados.py index aea5f5f..cd4efdd 100755 --- a/atrcopy/spartados.py +++ b/atrcopy/spartados.py @@ -179,19 +179,19 @@ class SpartaDosDiskImage(DiskImageBase): def get_boot_segments(self): segments = [] num = min(self.num_boot, 1) - bytes, style = self.get_sectors(1, num) + s = self.get_sector_slice(1, num) + r = self.rawdata[s] addr = self.boot_addr - header = ObjSegment(bytes[0:43], style[0:43], 0, 0, addr, addr + 43, name="Boot Header") + header = ObjSegment(r[0:43], 0, 0, addr, addr + 43, name="Boot Header") segments.append(header) if self.num_boot > 0: - sectors = ObjSegment(bytes, style, 0, 0, addr, addr + len(bytes), name="Boot Sectors") - code = ObjSegment(bytes[43:], style[43:], 0, 0, addr + 43, addr + len(bytes), name="Boot Code") + sectors = ObjSegment(r, 0, 0, addr, addr + len(r), name="Boot Sectors") + code = ObjSegment(r[43:], 0, 0, addr + 43, addr + len(r), name="Boot Code") segments.extend([header, code]) return segments def get_vtoc_segments(self): - b = self.bytes - s = self.style + r = self.rawdata segments = [] addr = 0 start, count = self.get_contiguous_sectors(self.first_bitmap, self.num_bitmap) @@ -201,7 +201,7 @@ class SpartaDosDiskImage(DiskImageBase): else: num_boot = 3 boot_size = 128 - segment = RawSectorsSegment(b[start:start+count], s[start:start+count], self.first_bitmap, self.num_bitmap, count, 0, 0, self.sector_size, name="Bitmap") + segment = RawSectorsSegment(r[start:start+count], self.first_bitmap, self.num_bitmap, count, 0, 0, self.sector_size, name="Bitmap") segments.append(segment) return segments @@ -236,7 +236,7 @@ class SpartaDosDiskImage(DiskImageBase): if len(byte_order) > 0: name = "%s %d@%d %s" % (dirent.get_filename(), dirent.length, dirent.starting_sector, dirent.str_timestamp) verbose_name = "%s (%d bytes, sector map@%d) %s %s" % (dirent.get_filename(), dirent.length, dirent.starting_sector, dirent.verbose_info, dirent.str_timestamp) - segment = IndexedByteSegment(self.bytes, self.style, byte_order, name=name, verbose_name=verbose_name) + segment = IndexedByteSegment(self.rawdata, byte_order, name=name, verbose_name=verbose_name) else: - segment = EmptySegment(self.bytes, self.style, name=dirent.get_filename(), error=dirent.str_timestamp) + segment = EmptySegment(self.rawdata, name=dirent.get_filename(), error=dirent.str_timestamp) return segment