From e4ce309bd368cd1899a5eb8f0a905b9b5fb29a92 Mon Sep 17 00:00:00 2001
From: Rob McMullen <feedback@playermissile.com>
Date: Tue, 17 Jul 2018 18:33:16 -0700
Subject: [PATCH] Added file extensions recognition to
 create_executable_file_image * you can write a XEX to an Apple DOS 3.3 image,
 for instance

---
 atrcopy/__init__.py    |  2 +-
 atrcopy/ataridos.py    |  6 ++----
 atrcopy/diskimages.py  | 11 +++++++++--
 atrcopy/dos33.py       |  6 ++----
 atrcopy/executables.py | 13 ++++++++++++-
 5 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/atrcopy/__init__.py b/atrcopy/__init__.py
index cfcfd22..0c1c90c 100644
--- a/atrcopy/__init__.py
+++ b/atrcopy/__init__.py
@@ -253,7 +253,7 @@ def assemble_segments(source_files, data_files, obj_files, run_addr=""):
 
 def assemble(image, source_files, data_files, obj_files, run_addr=""):
     segments, run_addr = assemble_segments(source_files, data_files, obj_files, run_addr)
-    file_data, filetype = image.create_executable_file_image(segments, run_addr)
+    file_data, filetype = image.create_executable_file_image(options.output, segments, run_addr)
     print("total file size: $%x (%d) bytes" % (len(file_data), len(file_data)))
     changed = save_file(image, options.output, filetype, file_data)
     if changed:
diff --git a/atrcopy/ataridos.py b/atrcopy/ataridos.py
index 06613b2..feaeeed 100644
--- a/atrcopy/ataridos.py
+++ b/atrcopy/ataridos.py
@@ -454,6 +454,8 @@ class XfdHeader(AtrHeader):
 
 
 class AtariDosDiskImage(DiskImageBase):
+    default_executable_extension = "XEX"
+
     def __init__(self, *args, **kwargs):
         self.first_vtoc = 360
         self.num_vtoc = 1
@@ -677,10 +679,6 @@ class AtariDosDiskImage(DiskImageBase):
                 log.debug("%s not a binary file; skipping segment generation" % str(segment))
         return segments_out
 
-    def create_executable_file_image(self, segments, run_addr=None):
-        base_segment, user_segments = get_xex(segments, run_addr)
-        return base_segment.data, "XEX"
-
 
 class BootDiskImage(AtariDosDiskImage):
     def __str__(self):
diff --git a/atrcopy/diskimages.py b/atrcopy/diskimages.py
index 6fd7551..3ac1328 100644
--- a/atrcopy/diskimages.py
+++ b/atrcopy/diskimages.py
@@ -3,6 +3,7 @@ import numpy as np
 from . import errors
 from .segments import SegmentData, EmptySegment, ObjSegment, RawSectorsSegment
 from .utils import *
+from .executables import create_executable_file_data
 
 import logging
 log = logging.getLogger(__name__)
@@ -92,6 +93,8 @@ class BaseHeader:
 
 
 class DiskImageBase:
+    default_executable_extension = None
+
     def __init__(self, rawdata, filename="", create=False):
         self.rawdata = rawdata
         self.bytes = self.rawdata.get_data()
@@ -300,8 +303,12 @@ class DiskImageBase:
             segments.append(segment)
         return segments
 
-    def create_executable_file_image(self, segments, run_addr=None):
-        raise errors.NotImplementedError
+    def create_executable_file_image(self, output_name, segments, run_addr=None):
+        try:
+            data, filetype = create_executable_file_data(output_name, segments, run_addr)
+        except errors.UnsupportedContainer:
+            data, filetype = create_executable_file_data(self.default_executable_extension, segments, run_addr)
+        return data, filetype
 
     @classmethod
     def create_boot_image(self, segments, run_addr=None):
diff --git a/atrcopy/dos33.py b/atrcopy/dos33.py
index 9fba871..8f1471c 100644
--- a/atrcopy/dos33.py
+++ b/atrcopy/dos33.py
@@ -382,6 +382,8 @@ class Dos33Header(BaseHeader):
 
 
 class Dos33DiskImage(DiskImageBase):
+    default_executable_extension = "BSAVE"
+
     def __init__(self, rawdata, filename=""):
         DiskImageBase.__init__(self, rawdata, filename)
         self.default_filetype = "B"
@@ -591,10 +593,6 @@ class Dos33DiskImage(DiskImageBase):
             segment = EmptySegment(self.rawdata, name=dirent.filename)
         return segment
 
-    def create_executable_file_image(self, segments, run_addr=None):
-        data = get_bsave(segments, run_addr)
-        return data, 'B'
-
 
 class Dos33BinFile:
     """Parse a binary chunk into segments according to the DOS 3.3 binary
diff --git a/atrcopy/executables.py b/atrcopy/executables.py
index d06f56f..dd52992 100644
--- a/atrcopy/executables.py
+++ b/atrcopy/executables.py
@@ -1,7 +1,6 @@
 import numpy as np
 
 from . import errors
-from .diskimages import DiskImageBase, BaseHeader
 from .segments import SegmentData, EmptySegment, ObjSegment, RawSectorsSegment, DefaultSegment, SegmentedFileSegment, SegmentSaver, get_style_bits
 from .utils import *
 
@@ -62,6 +61,7 @@ def get_xex(segments, run_addr=None):
         sub_segments.append(new_s)
     return main_segment, sub_segments
 
+
 def get_bsave(segments, run_addr=None):
     # Apple 2 executables get executed at the first address loaded. If the
     # run_addr is not the first byte of the combined data, have to create a
@@ -99,3 +99,14 @@ def get_bsave(segments, run_addr=None):
         print("setting data for $%04x - $%04x at index $%04x" % (s.origin, s.origin + len(s), index))
         image[index:index + len(s)] = s.data
     return image
+
+
+def create_executable_file_data(filename, segments, run_addr=None):
+    name = filename.lower()
+    if name.endswith("xex"):
+        base_segment, user_segments = get_xex(segments, run_addr)
+        return base_segment.data, "XEX"
+    elif name.endswith("bin") or name.endswith("bsave"):
+        data = get_bsave(segments, run_addr)
+        return data, "B"
+    raise errors.UnsupportedContainer