Spruce up make_shim.py

1 files changed, 61 insertions, 40 deletions

diff --git a/tools/make_shim.py b/tools/make_shim.py
index 7c60f4f..d5e6e00 100644
--- a/tools/make_shim.py
+++ b/tools/make_shim.py
@@ -1,4 +1,7 @@
-from sys import argv, stderr
+#!/usr/bin/env python
+
+from __future__ import print_function
+from sys import argv, stderr, exit
 
 
 section_list = [
@@ -20,15 +23,29 @@ file_list = []
 options = []
 while argv_id < len(argv):
     arg = argv[argv_id]
+
     if arg[0] != '-':
         file_list.append(arg)
+        argv_id += 1
+        continue
+
+    # An empty '--' stops parsing arguments
+    if arg == '--':
+        argv_id += 1
+        break
+
+    if arg[1] == '-':
+        options.append(option[2:])
     elif arg[1] != '-':
         for option in arg[1:]:
             options.append(option)
-    else:
-        options.append(option[2:])
+
     argv_id += 1
 
+# Add remaining files to the list
+for arg in argv[argv_id:]:
+    file_list.append(arg)
+
 
 if 'w' in options or 'd' in options:
     section_list[4]['end'] = 0xE000
@@ -38,40 +55,44 @@ if 't' in options:
 
 
 for file_name in file_list:
-    with open(file_name, "rt") as f:
-        for line in f:
-            line = line.split(";")[0]
-            line = line.strip()
-            if line and " " in line:
-                address, symbol = line.split()
-                bank, pointer = address.split(":")
-                bank = int(bank, 16)
-                pointer = int(pointer, 16)
-
-                section = None
-                for section_type in section_list:
-                    if pointer < section_type['end']:
-                        if section_type['invalid']:
-                            stderr.write(f"Warning: cannot shim '{symbol}' in section type '{section_type['name']}'\n")
-                            section = False
-                        else:
-                            section = section_type['name']
-                            if not section_type['banked']:
-                                bank = None
-                        break
-                
-                if section == False: # Found section, but cannot shim it
-                    continue
-                
-                if section == None: # Didn't find a section at all
-                    stderr.write(f"Unknown section for {line}\n")
-                
-                if section:
-                    if bank:
-                        print(f"SECTION \"Shim for {symbol}\", {section}[${pointer:x}], BANK[${bank:x}]")
-                    else:
-                        print(f"SECTION \"Shim for {symbol}\", {section}[${pointer:x}]")
-                    print(f"{symbol}::\n\n")
-
-    
-    
+    for line in open(file_name, "rt"):
+
+        # Strip out the comment
+        line = line.split(";")[0].strip()
+        if not line:
+            continue
+
+        # Read the address
+        try:
+            address, symbol = line.split()
+            bank, pointer = address.split(":")
+            bank = int(bank, 16)
+            pointer = int(pointer, 16)
+        except:
+            print("Error: Cannot parse line: %s" % line, file=stderr)
+            exit(1)
+
+        section = None
+        for section_type in section_list:
+            if pointer < section_type['end']:
+                if section_type['invalid']:
+                    print("Warning: cannot shim '%s' in section type '%s'" % (symbol, section_type['name']), file=stderr)
+                    section = False
+                else:
+                    section = section_type['name']
+                    if not section_type['banked']:
+                        bank = None
+                break
+        else:
+            # Didn't find a section
+            print("Unknown section for '%s'" % line, file=stderr)
+            continue
+
+        if not section:
+            # Found section, but cannot shim it
+            continue
+
+        print("SECTION \"Shim for %s\", %s[$%04X]" % (symbol, section, pointer), end='')
+        if bank:
+            print(", BANK[$%04X]" % bank, end='')
+        print("\n%s::\n\n" % symbol)