mirror of
https://github.com/robmcmullen/atrcopy.git
synced 2024-11-29 11:51:14 +00:00
487 lines
17 KiB
Python
487 lines
17 KiB
Python
from __future__ import print_function
|
|
from builtins import zip
|
|
from builtins import range
|
|
from builtins import object
|
|
import os
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from atrcopy import DefaultSegment, SegmentData, get_xex, interleave_segments, user_bit_mask, diff_bit_mask
|
|
from atrcopy.errors import *
|
|
from functools import reduce
|
|
|
|
|
|
def get_indexed(segment, num, scale):
|
|
indexes = np.arange(num) * scale
|
|
raw = segment.rawdata.get_indexed(indexes)
|
|
s = DefaultSegment(raw, segment.origin + indexes[0])
|
|
return s, indexes
|
|
|
|
class TestSegment1(object):
|
|
def setup(self):
|
|
self.segments = []
|
|
for i in range(8):
|
|
data = np.arange(1024, dtype=np.uint8) * i
|
|
r = SegmentData(data)
|
|
self.segments.append(DefaultSegment(r, i * 1024))
|
|
|
|
def test_xex(self):
|
|
items = [
|
|
[(0, 1, 2), 0],
|
|
]
|
|
|
|
for indexes, stuff in items:
|
|
s = [self.segments[i] for i in indexes]
|
|
s[1].style[0:500] = diff_bit_mask
|
|
s[1].set_comment_at(0, "comment 0")
|
|
s[1].set_comment_at(10, "comment 10")
|
|
s[1].set_comment_at(100, "comment 100")
|
|
print(list(s[1].iter_comments_in_segment()))
|
|
with pytest.raises(InvalidBinaryFile):
|
|
seg, subseg = get_xex(s, 0xbeef)
|
|
seg, subseg = get_xex(s)
|
|
assert tuple(seg.data[0:2]) == (0xff, 0xff)
|
|
# 2 bytes for the ffff
|
|
# 4 bytes per segment for start, end address
|
|
# An extra segment has been inserted for the run address!
|
|
size = reduce(lambda a, b:a + len(b), subseg, 0)
|
|
assert len(seg) == 2 + size
|
|
print(id(s[1]), list(s[1].iter_comments_in_segment()))
|
|
print(id(subseg[2]), list(subseg[2].iter_comments_in_segment()))
|
|
for i, c in s[1].iter_comments_in_segment():
|
|
assert c == subseg[2].get_comment(i + 4)
|
|
assert np.all(s[1].style[:] == subseg[2].style[4:])
|
|
|
|
def test_copy(self):
|
|
for s in self.segments:
|
|
d = s.rawdata
|
|
print("orig:", d.data.shape, d.is_indexed, d.data, id(d.data))
|
|
c = d.copy()
|
|
print("copy", c.data.shape, c.is_indexed, c.data, id(c.data))
|
|
assert c.data.shape == s.data.shape
|
|
assert id(c) != id(s)
|
|
assert np.all((c.data[:] - s.data[:]) == 0)
|
|
c.data[0:100] = 1
|
|
print(d.data)
|
|
print(c.data)
|
|
assert not np.all((c.data[:] - s.data[:]) == 0)
|
|
|
|
|
|
class TestIndexed(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.segment = DefaultSegment(r, 0)
|
|
|
|
def test_indexed(self):
|
|
assert not self.segment.rawdata.is_indexed
|
|
s, indexes = get_indexed(self.segment, 1024, 3)
|
|
assert s.rawdata.is_indexed
|
|
for i in range(len(indexes)):
|
|
assert s.get_raw_index(i) == indexes[i]
|
|
|
|
# get indexed into indexed, will result in every 9th byte
|
|
s2, indexes2 = get_indexed(s, 256, 3)
|
|
assert s2.rawdata.is_indexed
|
|
for i in range(len(indexes2)):
|
|
assert s2.get_raw_index(i) == indexes2[i] * 3
|
|
|
|
def test_indexed_sub(self):
|
|
base = self.segment
|
|
assert not base.rawdata.is_indexed
|
|
raw = base.rawdata[512:1536] # 1024 byte segment
|
|
sub = DefaultSegment(raw, 512)
|
|
|
|
assert not sub.rawdata.is_indexed
|
|
for i in range(len(sub)):
|
|
ri = sub.get_raw_index(i)
|
|
assert ri == sub.origin + i
|
|
assert sub[i] == base[ri]
|
|
start, end = sub.byte_bounds_offset()
|
|
assert start == 512
|
|
assert end == 1536
|
|
|
|
with pytest.raises(IndexError) as e:
|
|
# attempt to get indexes to 1024 * 3... Index to big => fail!
|
|
s, indexes = get_indexed(sub, 1024, 3)
|
|
|
|
# try with elements up to 256 * 3
|
|
s, indexes = get_indexed(sub, 256, 3)
|
|
print(sub.data)
|
|
print(indexes)
|
|
print(s.data[:])
|
|
assert s.rawdata.is_indexed
|
|
for i in range(len(indexes)):
|
|
ri = s.get_raw_index(i)
|
|
print(ri, "base[ri]=%d" % base[ri], i, indexes[i], "s[i]=%d" % s[i])
|
|
assert ri == sub.origin + indexes[i]
|
|
assert s[i] == base[ri]
|
|
start, end = s.byte_bounds_offset()
|
|
assert start == 0
|
|
assert end == len(base)
|
|
|
|
# get indexed into indexed, will result in every 9th byte
|
|
s2, indexes2 = get_indexed(s, 64, 3)
|
|
assert s2.rawdata.is_indexed
|
|
for i in range(len(indexes2)):
|
|
assert s2.get_raw_index(i) == sub.origin + indexes2[i] * 3
|
|
start, end = s.byte_bounds_offset()
|
|
assert start == 0
|
|
assert end == len(base)
|
|
|
|
def test_interleave(self):
|
|
base = self.segment
|
|
r1 = base.rawdata[512:1024] # 512 byte segment
|
|
s1 = DefaultSegment(r1, 512)
|
|
r2 = base.rawdata[1024:1536] # 512 byte segment
|
|
s2 = DefaultSegment(r2, 1024)
|
|
|
|
indexes1 = r1.get_indexes_from_base()
|
|
verify1 = np.arange(512, 1024, dtype=np.uint32)
|
|
assert np.array_equal(indexes1, verify1)
|
|
|
|
indexes2 = r2.get_indexes_from_base()
|
|
verify2 = np.arange(1024, 1536, dtype=np.uint32)
|
|
assert np.array_equal(indexes2, verify2)
|
|
|
|
s = interleave_segments([s1, s2], 2)
|
|
a = np.empty(len(s1) + len(s2), dtype=np.uint8)
|
|
a[0::4] = s1[0::2]
|
|
a[1::4] = s1[1::2]
|
|
a[2::4] = s2[0::2]
|
|
a[3::4] = s2[1::2]
|
|
print(list(s[:]))
|
|
print(list(a[:]))
|
|
print(s.rawdata.order)
|
|
assert np.array_equal(s[:], a)
|
|
|
|
s = interleave_segments([s1, s2], 4)
|
|
a = np.empty(len(s1) + len(s2), dtype=np.uint8)
|
|
a[0::8] = s1[0::4]
|
|
a[1::8] = s1[1::4]
|
|
a[2::8] = s1[2::4]
|
|
a[3::8] = s1[3::4]
|
|
a[4::8] = s2[0::4]
|
|
a[5::8] = s2[1::4]
|
|
a[6::8] = s2[2::4]
|
|
a[7::8] = s2[3::4]
|
|
assert np.array_equal(s[:], a)
|
|
|
|
def test_interleave_not_multiple(self):
|
|
base = self.segment
|
|
r1 = base.rawdata[512:1024] # 512 byte segment
|
|
s1 = DefaultSegment(r1, 512)
|
|
r2 = base.rawdata[1024:1536] # 512 byte segment
|
|
s2 = DefaultSegment(r2, 1024)
|
|
|
|
indexes1 = r1.get_indexes_from_base()
|
|
verify1 = np.arange(512, 1024, dtype=np.uint32)
|
|
assert np.array_equal(indexes1, verify1)
|
|
|
|
indexes2 = r2.get_indexes_from_base()
|
|
verify2 = np.arange(1024, 1536, dtype=np.uint32)
|
|
assert np.array_equal(indexes2, verify2)
|
|
|
|
s = interleave_segments([s1, s2], 3)
|
|
|
|
# when interleave size isn't a multiple of the length, the final array
|
|
# will reduce the size of the input array to force it to be a multiple.
|
|
size = (len(s1) // 3) * 3
|
|
assert len(s) == size * 2
|
|
a = np.empty(len(s), dtype=np.uint8)
|
|
a[0::6] = s1[0:size:3]
|
|
a[1::6] = s1[1:size:3]
|
|
a[2::6] = s1[2:size:3]
|
|
a[3::6] = s2[0:size:3]
|
|
a[4::6] = s2[1:size:3]
|
|
a[5::6] = s2[2:size:3]
|
|
assert np.array_equal(s[:], a)
|
|
|
|
def test_interleave_different_sizes(self):
|
|
base = self.segment
|
|
r1 = base.rawdata[512:768] # 256 byte segment
|
|
s1 = DefaultSegment(r1, 512)
|
|
r2 = base.rawdata[1024:1536] # 512 byte segment
|
|
s2 = DefaultSegment(r2, 1024)
|
|
|
|
indexes1 = r1.get_indexes_from_base()
|
|
verify1 = np.arange(512, 768, dtype=np.uint32)
|
|
assert np.array_equal(indexes1, verify1)
|
|
|
|
indexes2 = r2.get_indexes_from_base()
|
|
verify2 = np.arange(1024, 1536, dtype=np.uint32)
|
|
assert np.array_equal(indexes2, verify2)
|
|
|
|
s = interleave_segments([s1, s2], 3)
|
|
|
|
# when interleave size isn't a multiple of the length, the final array
|
|
# will reduce the size of the input array to force it to be a multiple.
|
|
size = (min(len(s1), len(s2)) // 3) * 3
|
|
assert size == (256 // 3) * 3
|
|
assert len(s) == size * 2
|
|
a = np.empty(len(s), dtype=np.uint8)
|
|
a[0::6] = s1[0:size:3]
|
|
a[1::6] = s1[1:size:3]
|
|
a[2::6] = s1[2:size:3]
|
|
a[3::6] = s2[0:size:3]
|
|
a[4::6] = s2[1:size:3]
|
|
a[5::6] = s2[2:size:3]
|
|
assert np.array_equal(s[:], a)
|
|
|
|
def test_copy(self):
|
|
s, indexes = get_indexed(self.segment, 1024, 3)
|
|
c = s.rawdata.copy()
|
|
print(c.data.shape, c.is_indexed)
|
|
print(id(c.data.np_data), id(s.data.np_data))
|
|
assert c.data.shape == s.data.shape
|
|
assert id(c) != id(s)
|
|
assert np.all((c.data[:] - s.data[:]) == 0)
|
|
c.data[0:100] = 1
|
|
assert not np.all((c.data[:] - s.data[:]) == 0)
|
|
|
|
|
|
class TestComments(object):
|
|
def setup(self):
|
|
data = np.ones([4000], dtype=np.uint8)
|
|
r = SegmentData(data)
|
|
self.segment = DefaultSegment(r, 0)
|
|
self.sub_segment = DefaultSegment(r[2:202], 2)
|
|
|
|
def test_locations(self):
|
|
s = self.segment
|
|
s.set_comment([[4,5]], "test1")
|
|
s.set_comment([[40,50]], "test2")
|
|
s.set_style_ranges([[2,100]], comment=True)
|
|
s.set_style_ranges([[200, 299]], data=True)
|
|
for i in range(1,4):
|
|
for j in range(1, 4):
|
|
# create some with overlapping regions, some without
|
|
r = [500*j, 500*j + 200*i + 200]
|
|
s.set_style_ranges([r], user=i)
|
|
s.set_user_data([r], i, i*10 + j)
|
|
r = [100, 200]
|
|
s.set_style_ranges([r], user=4)
|
|
s.set_user_data([r], 4, 99)
|
|
r = [3100, 3200]
|
|
s.set_style_ranges([r], user=4)
|
|
s.set_user_data([r], 4, 99)
|
|
|
|
s2 = self.sub_segment
|
|
print(len(s2))
|
|
copy = s2.get_comment_locations()
|
|
print(copy)
|
|
# comments at 4 and 40 in the original means 2 and 38 in the copy
|
|
orig = s.get_comment_locations()
|
|
assert copy[2] == orig[4]
|
|
assert copy[28] == orig[38]
|
|
|
|
def test_split_data_at_comment(self):
|
|
s = self.segment
|
|
s.set_style_ranges([[0,1000]], data=True)
|
|
for i in range(0, len(s), 25):
|
|
s.set_comment([[i,i+1]], "comment at %d" % i)
|
|
|
|
s2 = self.sub_segment
|
|
print(len(s2))
|
|
copy = s2.get_comment_locations()
|
|
print(copy)
|
|
# comments at 4 and 40 in the original means 2 and 38 in the copy
|
|
orig = s.get_comment_locations()
|
|
print(orig[0:200])
|
|
assert copy[2] == orig[4]
|
|
assert copy[28] == orig[38]
|
|
|
|
r = s2.get_entire_style_ranges([1], user=True)
|
|
print(r)
|
|
assert r == [((0, 23), 1), ((23, 48), 1), ((48, 73), 1), ((73, 98), 1), ((98, 123), 1), ((123, 148), 1), ((148, 173), 1), ((173, 198), 1), ((198, 200), 1)]
|
|
|
|
def test_split_data_at_comment2(self):
|
|
s = self.segment
|
|
start = 0
|
|
i = 0
|
|
for end in range(40, 1000, 40):
|
|
s.set_style_ranges([[start, end]], user=i)
|
|
start = end
|
|
i = (i + 1) % 8
|
|
for i in range(0, len(s), 25):
|
|
s.set_comment([[i,i+1]], "comment at %d" % i)
|
|
|
|
s2 = self.sub_segment
|
|
print(len(s2))
|
|
copy = s2.get_comment_locations()
|
|
print(copy)
|
|
# comments at 4 and 40 in the original means 2 and 38 in the copy
|
|
orig = s.get_comment_locations()
|
|
print(orig[0:200])
|
|
assert copy[2] == orig[4]
|
|
assert copy[28] == orig[38]
|
|
|
|
r = s2.get_entire_style_ranges([1], user=user_bit_mask)
|
|
print(r)
|
|
assert r == [((0, 38), 0), ((38, 48), 1), ((48, 73), 1), ((73, 78), 1), ((78, 118), 2), ((118, 158), 3), ((158, 198), 4), ((198, 200), 5)]
|
|
|
|
def test_restore_comments(self):
|
|
s = self.segment
|
|
s.set_style_ranges([[0,1000]], data=True)
|
|
for i in range(0, len(s), 5):
|
|
s.set_comment([[i,i+1]], "comment at %d" % i)
|
|
|
|
s1 = self.segment
|
|
print(len(s1))
|
|
indexes = [7,12]
|
|
r = s1.get_comment_restore_data([indexes])
|
|
print(r)
|
|
# force clear comments
|
|
s1.rawdata.extra.comments = {}
|
|
s1.style[indexes[0]:indexes[1]] = 0
|
|
r0 = s1.get_comment_restore_data([indexes])
|
|
print(r0)
|
|
for start, end, style, items in r0:
|
|
print(style)
|
|
assert np.all(style == 0)
|
|
for rawindex, comment in list(items.values()):
|
|
assert not comment
|
|
s1.restore_comments(r)
|
|
r1 = s1.get_comment_restore_data([indexes])
|
|
print(r1)
|
|
for item1, item2 in zip(r, r1):
|
|
print(item1)
|
|
print(item2)
|
|
for a1, a2 in zip(item1, item2):
|
|
print(a1, a2)
|
|
if hasattr(a1, "shape"):
|
|
assert np.all(a1 - a2 == 0)
|
|
else:
|
|
assert a1 == a2
|
|
|
|
s2 = self.sub_segment
|
|
print(len(s2))
|
|
indexes = [5,10]
|
|
r = s2.get_comment_restore_data([indexes])
|
|
print(r)
|
|
# force clear comments
|
|
s2.rawdata.extra.comments = {}
|
|
s2.style[indexes[0]:indexes[1]] = 0
|
|
r0 = s2.get_comment_restore_data([indexes])
|
|
print(r0)
|
|
for start, end, style, items in r0:
|
|
print(style)
|
|
assert np.all(style == 0)
|
|
for rawindex, comment in list(items.values()):
|
|
assert not comment
|
|
s2.restore_comments(r)
|
|
r2 = s2.get_comment_restore_data([indexes])
|
|
print(r2)
|
|
for item1, item2 in zip(r, r2):
|
|
print(item1)
|
|
print(item2)
|
|
for a1, a2 in zip(item1, item2):
|
|
print(a1, a2)
|
|
if hasattr(a1, "shape"):
|
|
assert np.all(a1 - a2 == 0)
|
|
else:
|
|
assert a1 == a2
|
|
|
|
for item1, item2 in zip(r1, r2):
|
|
print(item1)
|
|
print(item2)
|
|
# indexes won't be the same, but rawindexes and comments will
|
|
assert np.all(item1[2] - item2[2] == 0)
|
|
assert set(item1[3].values()) == set(item2[3].values())
|
|
|
|
|
|
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):
|
|
# check to see data a view of some rawdata will be the same when the
|
|
# rawdata is resized.
|
|
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
|
|
|
|
# Check that the small view has the same data as its parent
|
|
for i in range(offset):
|
|
assert s[i] == c[i + offset]
|
|
|
|
# keep a copy of the old raw data of the subset
|
|
oldraw = s.rawdata.copy()
|
|
oldid = id(s.rawdata)
|
|
|
|
requested = 8192
|
|
oldsize, newsize = c.resize(requested)
|
|
assert newsize == requested
|
|
s.replace_data(c) # s should point to the same offset in the resized data
|
|
assert id(s.rawdata) == oldid # segment rawdata object should be same
|
|
assert id(oldraw.order) == id(s.rawdata.order) # order the same
|
|
for i in range(offset): # check values compared to parent
|
|
assert s[i] == c[i + offset]
|
|
|
|
# check for changes in parent/view reflected so we see that it's
|
|
# pointing to the same array in memory
|
|
newbase = c.rawdata
|
|
newsub = s.rawdata
|
|
print(c.rawdata.data[offset:offset+offset])
|
|
print(s.rawdata.data[:])
|
|
s.rawdata.data[:] = 111
|
|
print(c.rawdata.data[offset:offset+offset])
|
|
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()
|
|
# t.test_indexed()
|
|
# t.test_indexed_sub()
|
|
# t.test_interleave()
|
|
# t = TestSegment1()
|
|
# t.setup()
|
|
# t.test_xex()
|
|
# t.test_copy()
|
|
# t = TestComments()
|
|
# t.setup()
|
|
# t.test_split_data_at_comment()
|
|
# t.test_restore_comments()
|
|
t = TestResize()
|
|
t.setup()
|
|
t.test_subset()
|