From 2dc8b77ddcea943c7f3c8b1f1bf00e0dfee0c4ff Mon Sep 17 00:00:00 2001
From: Fabian Dill <fabian.dill@web.de>
Date: Sun, 14 Nov 2021 21:03:17 +0100
Subject: [PATCH] Patch: consolidate some if trees

---
 Patch.py | 41 +++++++++++++++++------------------------
 1 file changed, 17 insertions(+), 24 deletions(-)

diff --git a/Patch.py b/Patch.py
index af0a5e5e..09f41277 100644
--- a/Patch.py
+++ b/Patch.py
@@ -47,19 +47,9 @@ def generate_yaml(patch: bytes, metadata: Optional[dict] = None, game: str = GAM
 
 
 def generate_patch(rom: bytes, metadata: Optional[dict] = None, game: str = GAME_ALTTP) -> bytes:
-    if game == GAME_ALTTP:
-        from worlds.alttp.Rom import get_base_rom_bytes
-    elif game == GAME_SM:
-        from worlds.sm.Rom import get_base_rom_bytes
-    elif game == GAME_SOE:
-        file_name = Utils.get_options()["soe_options"]["rom"]
-        get_base_rom_bytes = lambda: bytes(read_rom(open(file_name, "rb")))
-    else:
-        raise RuntimeError("Selected game for base rom not found.")
-
     if metadata is None:
         metadata = {}
-    patch = bsdiff4.diff(get_base_rom_bytes(), rom)
+    patch = bsdiff4.diff(get_base_rom_data(game), rom)
     return generate_yaml(patch, metadata, game)
 
 
@@ -80,27 +70,30 @@ def create_patch_file(rom_file_to_patch: str, server: str = "", destination: str
 def create_rom_bytes(patch_file: str, ignore_version: bool = False) -> Tuple[dict, str, bytearray]:
     data = Utils.parse_yaml(lzma.decompress(load_bytes(patch_file)).decode("utf-8-sig"))
     game_name = data["game"]
-    if game_name in supported_games:
-        if game_name == GAME_ALTTP:
-            from worlds.alttp.Rom import get_base_rom_bytes
-        elif game_name == GAME_SM:
-            from worlds.sm.Rom import get_base_rom_bytes
-        else:
-            raise Exception(f"No Patch handler for game {game_name}")
-    elif game_name == "alttp":  # old version for A Link to the Past
-        from worlds.alttp.Rom import get_base_rom_bytes
-    else:
-        raise Exception(f"Cannot handle game {game_name}")
-
     if not ignore_version and data["compatible_version"] > current_patch_version:
         raise RuntimeError("Patch file is incompatible with this patcher, likely an update is required.")
-    patched_data = bsdiff4.patch(get_base_rom_bytes(), data["patch"])
+    patched_data = bsdiff4.patch(get_base_rom_data(game_name), data["patch"])
     rom_hash = patched_data[int(0x7FC0):int(0x7FD5)]
     data["meta"]["hash"] = "".join(chr(x) for x in rom_hash)
     target = os.path.splitext(patch_file)[0] + ".sfc"
     return data["meta"], target, patched_data
 
 
+def get_base_rom_data(game: str):
+    if game == GAME_ALTTP:
+        from worlds.alttp.Rom import get_base_rom_bytes
+    elif game == "alttp":  # old version for A Link to the Past
+        from worlds.alttp.Rom import get_base_rom_bytes
+    elif game == GAME_SM:
+        from worlds.sm.Rom import get_base_rom_bytes
+    elif game == GAME_SOE:
+        file_name = Utils.get_options()["soe_options"]["rom"]
+        get_base_rom_bytes = lambda: bytes(read_rom(open(file_name, "rb")))
+    else:
+        raise RuntimeError("Selected game for base rom not found.")
+    return get_base_rom_bytes()
+
+
 def create_rom_file(patch_file: str) -> Tuple[dict, str]:
     data, target, patched_data = create_rom_bytes(patch_file)
     with open(target, "wb") as f: