Added resizable segments

This commit is contained in:
Rob McMullen 2017-03-21 16:25:59 -07:00
parent 7da17d65bc
commit 6f29e6053a
6 changed files with 174 additions and 10 deletions

View File

@ -8,10 +8,10 @@ except ImportError:
raise RuntimeError("atrcopy %s requires numpy" % __version__)
from errors import *
from ataridos import AtrHeader, AtariDosDiskImage, BootDiskImage, AtariDosFile, get_xex, add_atr_header
from ataridos import AtrHeader, AtariDosDiskImage, BootDiskImage, AtariDosFile, XexContainerSegment, get_xex, add_atr_header
from dos33 import Dos33DiskImage
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, SegmentList, get_style_mask, get_style_bits
from segments import SegmentData, SegmentSaver, DefaultSegment, EmptySegment, ObjSegment, RawSectorsSegment, SegmentedFileSegment, user_bit_mask, match_bit_mask, comment_bit_mask, data_style, selected_bit_mask, diff_bit_mask, not_user_bit_mask, interleave_segments, SegmentList, get_style_mask, get_style_bits
from spartados import SpartaDosDiskImage
from cartridge import A8CartHeader, AtariCartImage
from parsers import SegmentParser, DefaultSegmentParser, guess_parser_for_mime, guess_parser_for_system, iter_parsers, iter_known_segment_parsers, mime_parse_order

View File

@ -2,7 +2,7 @@ import numpy as np
from errors import *
from diskimages import DiskImageBase, BaseHeader
from segments import SegmentData, EmptySegment, ObjSegment, RawSectorsSegment, DefaultSegment, SegmentSaver, get_style_bits
from segments import SegmentData, EmptySegment, ObjSegment, RawSectorsSegment, DefaultSegment, SegmentedFileSegment, SegmentSaver, get_style_bits
from utils import *
import logging
@ -242,6 +242,10 @@ class XexSegmentSaver(SegmentSaver):
export_extensions = [".xex"]
class XexContainerSegment(DefaultSegment):
can_resize_default = True
class XexSegment(ObjSegment):
savers = [SegmentSaver, XexSegmentSaver]

View File

@ -2,7 +2,7 @@ import numpy as np
from segments import SegmentData, DefaultSegment
from kboot import KBootImage
from ataridos import AtariDosDiskImage, BootDiskImage, AtariDosFile
from ataridos import AtariDosDiskImage, BootDiskImage, AtariDosFile, XexContainerSegment
from spartados import SpartaDosDiskImage
from cartridge import AtariCartImage, get_known_carts
from mame import MameZipImage
@ -16,6 +16,7 @@ log = logging.getLogger(__name__)
class SegmentParser(object):
menu_name = ""
image_type = None
container_segment = DefaultSegment
def __init__(self, segment_data, strict=False):
self.image = None
@ -26,7 +27,7 @@ class SegmentParser(object):
def parse(self):
r = self.segment_data
self.segments.append(DefaultSegment(r, 0, name=self.menu_name))
self.segments.append(self.container_segment(r, 0, name=self.menu_name))
try:
self.image = self.get_image(r)
self.check_image()
@ -80,6 +81,7 @@ class AtariBootDiskSegmentParser(SegmentParser):
class XexSegmentParser(SegmentParser):
menu_name = "XEX (Atari 8-bit executable)"
image_type = AtariDosFile
container_segment = XexContainerSegment
class AtariCartSegmentParser(SegmentParser):

View File

@ -153,11 +153,45 @@ class SegmentData(object):
def __len__(self):
return len(self.data)
def resize(self, newsize):
if self.data.base is None:
try:
newdata = np.resize(self.data, (newsize,))
newstyle = np.resize(self.style, (newsize,))
except:
raise
else:
self.data = newdata
self.style = newstyle
def replace_arrays(self, base_raw):
newsize = len(base_raw)
oldsize = len(self.data_base)
if newsize < oldsize:
raise NotImplementedError("Can't truncate yet")
if self.is_indexed:
self.data.np_data = base_raw.data
self.data.base = base_raw.data.base
self.style.np_data = base_raw.style
self.style.base = base_raw.style.base
elif self.data.base is not None:
# if there is no base array, we aren't looking at a slice so we
# must be copying the entire array.
start, end = self.byte_bounds_offset()
self.data = base_raw.data[start:end]
self.style = base_raw.style[start:end]
else:
raise ValueError("The base SegmentData object should use the resize method to replace arrays")
@property
def stringio(self):
buf = cStringIO.StringIO(self.data[:])
return buf
@property
def is_base(self):
return not self.is_indexed and self.data.base is None
@property
def data_base(self):
return self.data.np_data if self.is_indexed else self.data.base if self.data.base is not None else self.data
@ -301,6 +335,8 @@ class SegmentData(object):
class DefaultSegment(object):
savers = [SegmentSaver]
use_origin_default = False
can_resize_default = False
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
@ -314,7 +350,11 @@ class DefaultSegment(object):
# Some segments may not have a standard place in memory, so this flag
# can be used to skip the memory map lookup when displaying disassembly
self.use_origin = False
self.use_origin = self.__class__.use_origin_default
# Some segments may be resized to contain additional segments not
# present when the segment was created.
self.can_resize = self.__class__.can_resize_default
def set_raw(self, rawdata):
self.rawdata = rawdata
@ -324,10 +364,43 @@ class DefaultSegment(object):
def get_raw(self):
return self.rawdata
def resize(self, newsize, zeros=True):
""" Resize the data arrays.
This can only be performed on the container segment. Child segments
must adjust their rawdata to point to the correct place.
Since segments don't keep references to other segments, it is the
user's responsibility to update any child segments that point to this
segment's data.
Numpy can't do an in-place resize on an array that has a view, so the
data must be replaced and all segments that point to that raw data must
also be changed. This has to happen outside this method because it
doesn't know the segment list of segments using itself as a base.
"""
if not self.can_resize:
raise ValueError("Segment %s can't be resized" % str(self))
# only makes sense for the container (outermost) object
if not self.rawdata.is_base:
raise ValueError("Only container segments can be resized")
origsize = len(self)
self.rawdata.resize(newsize)
self.set_raw(self.rawdata) # force attributes to be reset
newsize = len(self)
if zeros:
if newsize > origsize:
self.data[origsize:] = 0
self.style[origsize:] = 0
return origsize, newsize
def replace_data(self, container):
self.rawdata.replace_arrays(container.rawdata)
def __getstate__(self):
state = dict()
for key in ['start_addr', 'error', 'name', 'verbose_name', 'page_size', 'map_width', 'uuid']:
for key in ['start_addr', 'error', 'name', 'verbose_name', 'page_size', 'map_width', 'uuid', 'use_origin', 'can_resize']:
state[key] = getattr(self, key)
r = self.rawdata
state['_rawdata_bounds'] = list(r.byte_bounds_offset())
@ -347,6 +420,10 @@ class DefaultSegment(object):
"""
if not hasattr(self, 'uuid'):
self.uuid = str(uuid.uuid4())
if not hasattr(self, 'use_origin'):
self.use_origin = self.__class__.use_origin_default
if not hasattr(self, 'can_resize'):
self.can_resize = self.__class__.can_resize_default
def reconstruct_raw(self, rawdata):
start, end = self._rawdata_bounds
@ -918,6 +995,10 @@ class ObjSegment(DefaultSegment):
return s
class SegmentedFileSegment(ObjSegment):
can_resize_default = True
class RawSectorsSegment(DefaultSegment):
def __init__(self, rawdata, first_sector, num_sectors, count, boot_sector_size, num_boot_sectors, sector_size, **kwargs):
DefaultSegment.__init__(self, rawdata, 0, **kwargs)

View File

@ -1,6 +1,6 @@
from mock import *
from atrcopy import SegmentData, AtariDosFile, InvalidBinaryFile
from atrcopy import SegmentData, AtariDosFile, InvalidBinaryFile, DefaultSegment, XexContainerSegment
class TestAtariDosFile(object):
@ -8,12 +8,24 @@ class TestAtariDosFile(object):
pass
def test_segment(self):
bytes = [0xff, 0xff, 0x00, 0x60, 0x01, 0x60, 1, 2]
bytes = np.asarray([0xff, 0xff, 0x00, 0x60, 0x01, 0x60, 1, 2], dtype=np.uint8)
rawdata = SegmentData(bytes)
image = AtariDosFile(rawdata)
container = XexContainerSegment(rawdata, 0)
image = AtariDosFile(container.rawdata)
image.parse_segments()
print image.segments
assert len(image.segments) == 1
assert len(image.segments[0]) == 2
assert np.all(image.segments[0] == bytes[6:8])
container.resize(16)
for s in image.segments:
s.replace_data(container)
new_segment = DefaultSegment(rawdata[8:16])
new_segment[:] = 99
assert np.all(image.segments[0] == bytes[6:8])
print new_segment[:]
assert np.all(new_segment[:] == 99)
def test_short_segment(self):
bytes = [0xff, 0xff, 0x00, 0x60, 0xff, 0x60, 1, 2]

View File

@ -331,6 +331,71 @@ class TestComments(object):
assert item1[3] == item2[3]
class TestResize(object):
def setup(self):
data = np.arange(4096, dtype=np.uint8)
data[1::2] = np.repeat(np.arange(16, dtype=np.uint8), 128)
r = SegmentData(data)
self.container = DefaultSegment(r, 0)
self.container.can_resize = True
def test_subset(self):
c = self.container
assert not c.rawdata.is_indexed
offset = 1000
s = DefaultSegment(c.rawdata[offset:offset + offset], 0)
assert not s.rawdata.is_indexed
for i in range(offset):
assert s[i] == c[i + offset]
requested = 8192
oldraw = s.rawdata.copy()
oldid = id(s.rawdata)
oldsize, newsize = c.resize(requested)
assert newsize == requested
s.replace_data(c)
assert id(s.rawdata) == oldid
assert id(oldraw.order) == id(s.rawdata.order)
for i in range(offset):
assert s[i] == c[i + offset]
newbase = c.rawdata
newsub = s.rawdata
print c.rawdata.data
print s.rawdata.data[:]
s.rawdata.data[:] = 111
print c.rawdata.data
print s.rawdata.data[:]
for i in range(offset):
assert s[i] == c[i + offset]
def test_indexed(self):
c = self.container
assert not c.rawdata.is_indexed
s, indexes = get_indexed(self.container, 1024, 3)
assert s.rawdata.is_indexed
for i in range(len(indexes)):
assert s.get_raw_index(i) == indexes[i]
requested = 8192
oldraw = s.rawdata.copy()
oldid = id(s.rawdata)
oldsize, newsize = c.resize(requested)
assert newsize == requested
s.replace_data(c)
assert id(s.rawdata) == oldid
assert id(oldraw.order) == id(s.rawdata.order)
for i in range(len(indexes)):
assert s.get_raw_index(i) == indexes[i]
newbase = c.rawdata
newsub = s.rawdata
print c.rawdata.data
print s.rawdata.data[:]
s.rawdata.data[:] = 111
print c.rawdata.data
print s.rawdata.data[:]
for i in range(len(indexes)):
assert c.rawdata.data[indexes[i]] == s.rawdata.data[i]
if __name__ == "__main__":
t = TestIndexed()
t.setup()