From 3e77cb86bcb501a4b6cf0aee29b795a5bf68a854 Mon Sep 17 00:00:00 2001
From: Rob McMullen <feedback@playermissile.com>
Date: Fri, 24 Feb 2017 23:48:10 -0800
Subject: [PATCH] Added automatic data byte marking for XEX expanded files

---
 atrcopy/ataridos.py |  6 +++-
 atrcopy/segments.py | 72 ++++++++++++++++++++++++---------------------
 2 files changed, 44 insertions(+), 34 deletions(-)

diff --git a/atrcopy/ataridos.py b/atrcopy/ataridos.py
index 15bb1b8..abd6708 100644
--- a/atrcopy/ataridos.py
+++ b/atrcopy/ataridos.py
@@ -2,7 +2,7 @@ import numpy as np
 
 from errors import *
 from diskimages import DiskImageBase, BaseHeader
-from segments import EmptySegment, ObjSegment, RawSectorsSegment, DefaultSegment, SegmentSaver
+from segments import EmptySegment, ObjSegment, RawSectorsSegment, DefaultSegment, SegmentSaver, get_style_bits
 from utils import *
 
 import logging
@@ -264,7 +264,9 @@ class AtariDosFile(object):
     def parse_segments(self):
         r = self.rawdata
         b = r.get_data()
+        s = r.get_style()
         pos = 0
+        style_pos = 0
         first = True
         log.debug("Initial parsing: size=%d" % self.size)
         while pos < self.size:
@@ -286,6 +288,7 @@ class AtariDosFile(object):
                 self.segments.append(ObjSegment(r[pos:pos + 4], 0, 0, 0, len(b[pos:pos + 4]), "Short Segment Header"))
                 break
             start, end = b[pos:pos + 4].view(dtype='<u2')
+            s[style_pos:pos + 4] = get_style_bits(data=True)
             if end < start:
                 raise InvalidBinaryFile("Nonsensical start and end addresses")
             count = end - start + 1
@@ -295,6 +298,7 @@ class AtariDosFile(object):
                 break
             self.segments.append(ObjSegment(r[pos + 4:pos + 4 + count], pos, pos + 4, start, end))
             pos += 4 + count
+            style_pos = pos
 
 
 class AtrHeader(BaseHeader):
diff --git a/atrcopy/segments.py b/atrcopy/segments.py
index 6a6975a..5ea2102 100644
--- a/atrcopy/segments.py
+++ b/atrcopy/segments.py
@@ -15,6 +15,42 @@ match_bit_mask = 0x20
 comment_bit_mask = 0x40
 selected_bit_mask = 0x80
 
+def get_style_bits(match=False, comment=False, selected=False, data=False, diff=False, user=0):
+    """ Return an int value that contains the specified style bits set.
+    
+    Available styles for each byte are:
+    
+    match: part of the currently matched search
+    comment: user commented area
+    selected: selected region
+    data: labeled in the disassembler as a data region (i.e. not disassembled)
+    """
+    style_bits = 0
+    if user:
+        style_bits |= (user & user_bit_mask)
+    if diff:
+        style_bits |= diff_bit_mask
+    if match:
+        style_bits |= match_bit_mask
+    if comment:
+        style_bits |= comment_bit_mask
+    if data:
+        style_bits |= (data_style & user_bit_mask)
+    if selected:
+        style_bits |= selected_bit_mask
+    return style_bits
+
+def get_style_mask(**kwargs):
+    """Get the bit mask that, when anded with data, will turn off the
+    selected bits
+    """
+    bits = get_style_bits(**kwargs)
+    if 'user' in kwargs and kwargs['user']:
+        bits |= user_bit_mask
+    else:
+        bits &= (0xff ^ user_bit_mask)
+    return 0xff ^ bits
+    
 
 class SegmentSaver(object):
     export_data_name = "Raw Data"
@@ -399,41 +435,11 @@ class DefaultSegment(object):
     def tostring(self):
         return self.data.tostring()
     
-    def get_style_bits(self, match=False, comment=False, selected=False, data=False, diff=False, user=0):
-        """ Return an int value that contains the specified style bits set.
-        
-        Available styles for each byte are:
-        
-        match: part of the currently matched search
-        comment: user commented area
-        selected: selected region
-        data: labeled in the disassembler as a data region (i.e. not disassembled)
-        """
-        style_bits = 0
-        if user:
-            style_bits |= (user & user_bit_mask)
-        if diff:
-            style_bits |= diff_bit_mask
-        if match:
-            style_bits |= match_bit_mask
-        if comment:
-            style_bits |= comment_bit_mask
-        if data:
-            style_bits |= (data_style & user_bit_mask)
-        if selected:
-            style_bits |= selected_bit_mask
-        return style_bits
+    def get_style_bits(self, **kwargs):
+        return get_style_bits(**kwargs)
     
     def get_style_mask(self, **kwargs):
-        """Get the bit mask that, when anded with data, will turn off the
-        selected bits
-        """
-        bits = self.get_style_bits(**kwargs)
-        if 'user' in kwargs and kwargs['user']:
-            bits |= user_bit_mask
-        else:
-            bits &= (0xff ^ user_bit_mask)
-        return 0xff ^ bits
+        return get_style_mask(**kwargs)
     
     def set_style_ranges(self, ranges, **kwargs):
         style_bits = self.get_style_bits(**kwargs)