diff --git a/README.rst b/README.rst index 928cf9a..80bb6a6 100644 --- a/README.rst +++ b/README.rst @@ -4,6 +4,12 @@ ATRCopy Utilities to list files on and extract files from Atari 8-bit emulator disk images. Eventually, I hope to add support for these images to pyfilesystem. +Prerequisites +------------- + +* Requires numpy + + References ========== diff --git a/atrcopy.py b/atrcopy.py index 5d2a5f7..47021ad 100755 --- a/atrcopy.py +++ b/atrcopy.py @@ -1,11 +1,12 @@ #!/usr/bin/env python -__version__ = "1.3.0" +__version__ = "2.0.0" +import types + +import numpy as np -import struct -from cStringIO import StringIO class AtrError(RuntimeError): pass @@ -23,7 +24,16 @@ class ByteNotInFile166(AtrError): pass class AtrHeader(object): - format = " 0 @@ -117,10 +134,10 @@ class AtrDirent(object): self.locked = (flag&0x20) > 0 self.in_use = (flag&0x40) > 0 self.deleted = (flag&0x80) > 0 - self.num_sectors = values[1] - self.starting_sector = values[2] - self.filename = values[3].rstrip() - self.ext = values[4].rstrip() + self.num_sectors = int(values[1]) + self.starting_sector = int(values[2]) + self.filename = str(values[3]).rstrip() + self.ext = str(values[4]).rstrip() self.current_sector = 0 self.is_sane = self.sanity_check(disk) @@ -150,17 +167,29 @@ class AtrDirent(object): self.current_read = self.num_sectors def read_sector(self, disk): - raw = disk.get_raw_bytes(self.current_sector) - bytes = self.process_raw_sector(disk, raw) - return (bytes, self.current_sector == 0) + raw, pos, size = disk.get_raw_bytes(self.current_sector) + bytes, num_data_bytes = self.process_raw_sector(disk, raw) + return bytes, self.current_sector == 0, pos, num_data_bytes def process_raw_sector(self, disk, raw): - file_num = ord(raw[-3]) >> 2 + try: + file_num = ord(raw[-3]) >> 2 + except TypeError: + # if numpy data, don't need the ord() + return self.process_raw_sector_numpy(disk, raw) if file_num != self.file_num: raise FileNumberMismatchError164() self.current_sector = ((ord(raw[-3]) & 0x3) << 8) + ord(raw[-2]) num_bytes = ord(raw[-1]) - return raw[0:num_bytes] + return raw[0:num_bytes], num_bytes + + def process_raw_sector_numpy(self, disk, raw): + file_num = raw[-3] >> 2 + if file_num != self.file_num: + raise FileNumberMismatchError164() + self.current_sector = ((raw[-3] & 0x3) << 8) + raw[-2] + num_bytes = raw[-1] + return raw[0:num_bytes], num_bytes def get_filename(self): ext = ("." + self.ext) if self.ext else "" @@ -181,38 +210,97 @@ class MydosDirent(AtrDirent): class InvalidBinaryFile(AtrError): pass -class ObjSegment(object): - def __init__(self, metadata_start, data_start, start_addr, end_addr, data, name="", error=None): - self.name = name - self.metadata_start = metadata_start - self.data_start = data_start - self.start_addr = start_addr - self.end_addr = end_addr + + +class DefaultSegment(object): + debug = False + + def __init__(self, start_addr=0, data=None, name="All", error=None): + self.start_addr = int(start_addr) # force python int to decouple from possibly being a numpy datatype + if data is None: + data = np.fromstring("", dtype=np.uint8) + else: + data = to_numpy(data) self.data = data + self.style = np.zeros_like(self.data, dtype=np.uint8) + if self.debug: + self.style = np.arange(len(self), dtype=np.uint8) self.error = error - if name and not name.endswith(" "): - name += " " self.name = name self.page_size = -1 + self.map_width = 40 + self._search_copy = None def __str__(self): - s = "%s%04x-%04x (%04x @ %04x)" % (self.name, self.start_addr, self.end_addr, len(self.data), self.data_start) + return "%s (%d bytes)" % (self.name, len(self.data)) + + def __len__(self): + return np.alen(self.data) + + def __getitem__(self, index): + return self.data[index] + + def __setitem__(self, index, value): + self.data[index] = value + self._search_copy = None + + def tostring(self): + return self.data.tostring() + + def get_style_bits(self, match=False, comment=False): + style_bits = 0 + if match: + style_bits |= 1 + if comment: + style_bits |= 0x80 + return style_bits + + def get_style_mask(self, match=False, comment=False): + style_mask = 0xff + if match: + style_mask &= 0xfe + if comment: + style_mask &= 0x7f + return style_mask + + def set_style_ranges(self, ranges, **kwargs): + style_bits = self.get_style_bits(**kwargs) + s = self.style + for start, end in ranges: + s[start:end] |= style_bits + + def clear_style_bits(self, **kwargs): + style_mask = self.get_style_mask(**kwargs) + self.style &= style_mask + + def label(self, index, lower_case=True): + if lower_case: + return "%04x" % (index + self.start_addr) + else: + return "%04X" % (index + self.start_addr) + + @property + def search_copy(self): + if self._search_copy is None: + self._search_copy = self.data.tostring() + return self._search_copy + +class ObjSegment(DefaultSegment): + def __init__(self, metadata_start, data_start, start_addr, end_addr, data, name="", error=None): + DefaultSegment.__init__(self, start_addr, data, name, error) + self.metadata_start = metadata_start + self.data_start = data_start + + def __str__(self): + count = len(self) + s = "%s%04x-%04x (%04x @ %04x)" % (self.name, self.start_addr, self.start_addr + count, count, self.data_start) if self.error: s += " " + self.error return s - - def __len__(self): - return len(self.data) - - def __getitem__(self, val): - return self.data[val] - - def label(self, index): - return "%04x" % (index + self.start_addr) -class RawSectorsSegment(ObjSegment): +class RawSectorsSegment(DefaultSegment): def __init__(self, first_sector, num_sectors, count, data, **kwargs): - ObjSegment.__init__(self, 0, 0, 0, count, data, **kwargs) + DefaultSegment.__init__(self, 0, data, **kwargs) self.page_size = 128 self.first_sector = first_sector self.num_sectors = num_sectors @@ -226,9 +314,17 @@ class RawSectorsSegment(ObjSegment): s += " " + self.error return s - def label(self, index): + def label(self, index, lower_case=True): sector, byte = divmod(index, self.page_size) - return "s%03d:%02x" % (sector + self.first_sector, byte) + if lower_case: + return "s%03d:%02x" % (sector + self.first_sector, byte) + return "s%03d:%02X" % (sector + self.first_sector, byte) + +class IndexedByteSegment(DefaultSegment): + def __init__(self, byte_order, bytes, **kwargs): + data = bytes[byte_order] + DefaultSegment.__init__(self, 0, data, **kwargs) + class AtariDosFile(object): """Parse a binary chunk into segments according to the Atari DOS object @@ -237,8 +333,8 @@ class AtariDosFile(object): Ref: http://www.atarimax.com/jindroush.atari.org/afmtexe.html """ def __init__(self, data): - self.data = data - self.size = len(data) + self.data = to_numpy(data) + self.size = len(self.data) self.segments = [] self.parse_segments() @@ -257,7 +353,7 @@ class AtariDosFile(object): pos = 0 first = True while pos < self.size: - header, = struct.unpack(" 0: self.segments.append(self.get_obj_segment(0, 0, 0, self.header.atr_header_offset, self.bytes[0:self.header.atr_header_offset], name="%s Header" % self.header.file_format)) @@ -466,6 +602,7 @@ class AtrDiskImage(object): self.segments.extend(self.get_boot_segments()) self.segments.extend(self.get_vtoc_segments()) self.segments.extend(self.get_directory_segments()) + self.segments.extend(self.get_file_segments()) # for dirent in self.atr.files: # try: @@ -480,6 +617,13 @@ class AtrDiskImage(object): # a = AtrFileSegment(dirent, bytes, error) # self.segments.append(AtrSegment(dirent)) +def to_numpy(value): + if type(value) is np.ndarray: + return value + elif type(value) is types.StringType: + return np.fromstring(value, dtype=np.uint8) + raise TypeError("Can't convert to numpy data") + def process(dirent, options): skip = False action = "copying to" diff --git a/setup.py b/setup.py index 6a9b59d..58e1549 100644 --- a/setup.py +++ b/setup.py @@ -28,4 +28,7 @@ setup(name="atrcopy", long_description=long_description, license="GPL", classifiers=classifiers, + install_requires = [ + 'numpy', + ], )