simplify load_rom in gbz80disasm

2 files changed, 7 insertions, 8 deletions

diff --git a/extras/gbz80disasm.py b/extras/gbz80disasm.py
index 088a32f74..d22f152f1 100644
--- a/extras/gbz80disasm.py
+++ b/extras/gbz80disasm.py
@@ -14,11 +14,8 @@ if not hasattr(json, "read"):
 from romstr import RomStr
 
 def load_rom(filename="../baserom.gbc"):
-    """loads bytes into memory"""
     global rom
-    file_handler = open(filename, "rb")
-    rom = RomStr(file_handler.read())
-    file_handler.close()
+    rom = RomStr.load(filename=filename)
     return rom
 
 spacing = "\t"
diff --git a/extras/romstr.py b/extras/romstr.py
index d2eea44ae..5701f19ae 100644
--- a/extras/romstr.py
+++ b/extras/romstr.py
@@ -46,15 +46,17 @@ class RomStr(str):
         return "RomStr(too long)"
 
     @classmethod
-    def load(cls, crystal=True, red=False):
+    def load(cls, filename=None, crystal=True, red=False):
         """ Loads a ROM into a RomStr.
         """
-        if   crystal and not red:
+        if   crystal and not red and not filename:
             file_handler = open("../baserom.gbc", "r")
-        elif red and not crystal:
+        elif red and not crystal and not filename:
             file_handler = open("../pokered-baserom.gbc", "r")
+        elif filename not in ["", None]:
+            file_handler = open(filename, "rb")
         else:
-            raise Exception, "not sure which rom to load?"
+            raise Exception("not sure which rom to load?")
         bytes = file_handler.read()
         file_handler.close()
         return RomStr(bytes)