atrcopy/test/test_segment.py

import os

import numpy as np
import pytest

from atrcopy import DefaultSegment, SegmentData, get_xex, interleave_segments, user_bit_mask, diff_bit_mask


def get_indexed(segment, num, scale):
    indexes = np.arange(num) * scale
    raw = segment.rawdata.get_indexed(indexes)
    s = DefaultSegment(raw, segment.start_addr + 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())
            seg, subseg = get_xex(s, 0xbeef)
            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.start_addr + 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.start_addr + 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.start_addr + 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)
        
        with pytest.raises(ValueError) as e:
            s = interleave_segments([s1, s2], 3)

        r1 = base.rawdata[512:1025]  # 513 byte segment
        s1 = DefaultSegment(r1, 512)
        r2 = base.rawdata[1024:1537]  # 513 byte segment
        s2 = DefaultSegment(r2, 1024)
        s = interleave_segments([s1, s2], 3)
        a = np.empty(len(s1) + len(s2), dtype=np.uint8)
        a[0::6] = s1[0::3]
        a[1::6] = s1[1::3]
        a[2::6] = s1[2::3]
        a[3::6] = s2[0::3]
        a[4::6] = s2[1::3]
        a[5::6] = s2[2::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 items:
                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 items:
                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 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()
    # 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()