changeset 298:92a6fd2632f1

Improve heuristic to filter out non-game exes (like custom.exe).
author Thibaut Girka <thib@sitedethib.com>
date Tue, 06 Mar 2012 18:58:41 +0100
parents a09ac4650e0d
children e04e402e6380
files pytouhou/formats/exe.py pytouhou/resource/loader.py
diffstat 2 files changed, 25 insertions(+), 8 deletions(-) [+]
line wrap: on
line diff
--- a/pytouhou/formats/exe.py
+++ b/pytouhou/formats/exe.py
@@ -26,6 +26,10 @@ logger = get_logger(__name__)
 SQ2 = 2. ** 0.5 / 2.
 
 
+class InvalidExeException(Exception):
+    pass
+
+
 class Shot(object):
     def __init__(self):
         self.interval = 0
@@ -141,7 +145,11 @@ class SHT(object):
         text_va = pe_file.image_base + text_section.VirtualAddress
         text_size = text_section.SizeOfRawData
 
-        character_records_va = list(cls.find_character_defs(pe_file))[0]
+        possible_character_records = list(cls.find_character_defs(pe_file))
+        if not possible_character_records:
+            raise InvalidExeException
+
+        character_records_va = possible_character_records[0]
 
         characters = []
         shots_offsets = {}
--- a/pytouhou/resource/loader.py
+++ b/pytouhou/resource/loader.py
@@ -23,11 +23,16 @@ from pytouhou.formats.ecl import ECL
 from pytouhou.formats.anm0 import ANM0
 from pytouhou.formats.msg import MSG
 from pytouhou.formats.sht import SHT
-from pytouhou.formats.exe import SHT as EoSDSHT
+from pytouhou.formats.exe import SHT as EoSDSHT, InvalidExeException
 
 
 from pytouhou.resource.anmwrapper import AnmWrapper
 
+from pytouhou.utils.helpers import get_logger
+
+logger = get_logger(__name__)
+
+
 
 class Directory(object):
     def __init__(self, path):
@@ -93,7 +98,7 @@ class ArchiveDescription(object):
 
 class Loader(object):
     def __init__(self, game_dir=None):
-        self.exe = None
+        self.exe_files = []
         self.game_dir = game_dir
         self.known_files = {}
         self.instanced_ecls = {}
@@ -113,7 +118,7 @@ class Loader(object):
             paths = list(chain(*_expand_paths()))
             path = paths[0]
             if os.path.splitext(path)[1] == '.exe':
-                self.exe = path
+                self.exe_files.extend(paths)
             else:
                 archive_description = ArchiveDescription.get_from_path(path)
                 for name in archive_description.file_list:
@@ -169,10 +174,14 @@ class Loader(object):
 
     def get_eosd_characters(self):
         #TODO: Move to pytouhou.games.eosd?
-        path = self.exe
-        with open(path, 'rb') as file:
-            characters = EoSDSHT.read(file) #TODO: modular
-        return characters
+        for path in self.exe_files:
+            try:
+                with open(path, 'rb') as file:
+                    characters = EoSDSHT.read(file)
+                return characters
+            except InvalidExeException:
+                pass
+        logger.error("Required game exe not found!")
 
 
     def get_anm_wrapper(self, names, offsets=None):