From 053c7a773ee18915e5434cf1d742ad604ff1c69f Mon Sep 17 00:00:00 2001
From: Rob McMullen <feedback@playermissile.com>
Date: Sun, 26 Mar 2017 15:41:07 -0700
Subject: [PATCH] Changed get_xex to return segments: root segment & list of
 sub-segments

---
 atrcopy/ataridos.py  | 29 +++++++++++++++++++++--------
 atrcopy/segments.py  |  9 +++++++++
 test/test_segment.py | 41 +++++++++++++++++++++++++----------------
 3 files changed, 55 insertions(+), 24 deletions(-)

diff --git a/atrcopy/ataridos.py b/atrcopy/ataridos.py
index d248c30..58ce777 100644
--- a/atrcopy/ataridos.py
+++ b/atrcopy/ataridos.py
@@ -720,6 +720,10 @@ class BootDiskImage(AtariDosDiskImage):
 
 
 def get_xex(segments, runaddr=None):
+    segments_copy = [s for s in segments]  # don't affect the original list!
+    main_segment = None
+    sub_segments = []
+    data_style = get_style_bits(data=True)
     total = 2
     runad = False
     for s in segments:
@@ -733,19 +737,28 @@ def get_xex(segments, runaddr=None):
         words[0] = runaddr
         r = SegmentData(words.view(dtype=np.uint8))
         s = DefaultSegment(r, 0x2e0)
-        segments[0:0] = [s]
+        segments_copy[0:0] = [s]
         total += 6
     bytes = np.zeros([total], dtype=np.uint8)
-    bytes[0:2] = 0xff # FFFF header
+    rawdata = SegmentData(bytes)
+    main_segment = DefaultSegment(rawdata)
+    main_segment.data[0:2] = 0xff  # FFFF header
+    main_segment.style[0:2] = data_style
     i = 2
-    for s in segments:
-        words = bytes[i:i+4].view(dtype='<u2')
+    for s in segments_copy:
+        # create new sub-segment inside new main segment that duplicates the
+        # original segment's data/style
+        new_s = DefaultSegment(rawdata[i:i+4+len(s)], s.start_addr)
+        words = new_s.data[0:4].view(dtype='<u2')
         words[0] = s.start_addr
         words[1] = s.start_addr + len(s) - 1
-        i += 4
-        bytes[i:i + len(s)] = s[:]
-        i += len(s)
-    return bytes
+        new_s.style[0:4] = data_style
+        new_s.data[4:4+len(s)] = s[:]
+        new_s.style[4:4+len(s)] = s.style[:]
+        i += 4 + len(s)
+        new_s.copy_user_data(s, 4)
+        sub_segments.append(new_s)
+    return main_segment, sub_segments
 
 
 def add_atr_header(bytes):
diff --git a/atrcopy/segments.py b/atrcopy/segments.py
index d399752..0f3f0b7 100644
--- a/atrcopy/segments.py
+++ b/atrcopy/segments.py
@@ -924,6 +924,15 @@ class DefaultSegment(object):
             if k >= start_index and k < end_index:
                 yield self.rawdata.get_reverse_index(k), v
 
+    def copy_user_data(self, source, index_offset=0):
+        """Copy comments and other user data from the source segment to this
+        segment.
+
+        The index offset is the offset into self based on the index of source.
+        """
+        for index, comment in source.iter_comments_in_segment():
+            self.set_comment_at(index + index_offset, comment)
+
     def label(self, index, lower_case=True):
         if lower_case:
             return "%04x" % (index + self.start_addr)
diff --git a/test/test_segment.py b/test/test_segment.py
index 884aa3a..58eb68e 100644
--- a/test/test_segment.py
+++ b/test/test_segment.py
@@ -3,7 +3,7 @@ import os
 import numpy as np
 import pytest
 
-from atrcopy import DefaultSegment, SegmentData, get_xex, interleave_segments, user_bit_mask
+from atrcopy import DefaultSegment, SegmentData, get_xex, interleave_segments, user_bit_mask, diff_bit_mask
 
 
 def get_indexed(segment, num, scale):
@@ -27,14 +27,23 @@ class TestSegment1(object):
         
         for indexes, stuff in items:
             s = [self.segments[i] for i in indexes]
-            bytes = get_xex(s, 0xbeef)
-            assert tuple(bytes[0:2]) == (0xff, 0xff)
+            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 + 4 + len(b), s, 0)
-            print size, len(bytes)
-            assert len(bytes) == 2 + size
+            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:
@@ -397,16 +406,16 @@ class TestResize(object):
 
 
 if __name__ == "__main__":
-    t = TestIndexed()
-    t.setup()
-    t.test_indexed()
-    t.test_indexed_sub()
-    t.test_interleave()
+    # 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.test_copy()
+    # t = TestComments()
+    # t.setup()
+    # t.test_split_data_at_comment()
+    # t.test_restore_comments()