From 4d8760c733d62d78b9adb21fa71b30d272708677 Mon Sep 17 00:00:00 2001
From: Joseph Montanaro <jfmonty2@gmail.com>
Date: Tue, 20 Apr 2021 20:40:04 -0700
Subject: [PATCH] improve performance by selectively extracting textures

---
 aphremove/__init__.py | 81 +++++++++++++++++++++----------------------
 setup.py              |  3 +-
 2 files changed, 42 insertions(+), 42 deletions(-)

diff --git a/aphremove/__init__.py b/aphremove/__init__.py
index 6b974d0..9b50987 100644
--- a/aphremove/__init__.py
+++ b/aphremove/__init__.py
@@ -1,64 +1,63 @@
 import json
-import os, os.path
+import os
 import pathlib
 import sys
 import tempfile
 
-from deppth import deppth # not sure why the nested path
+from deppth import deppth, sggpio # not sure why the nested path
 from PIL import Image, ImageDraw
 
 
-def get_list(filename):
-    d = pathlib.Path(__file__).parent
-    lines = (d / filename).read_text().splitlines()
-    return set(lines)
+entries_file = pathlib.Path(__file__).parent / 'entries.json'
+with open(entries_file) as f:
+    ENTRIES = json.load(f)
 
 
-def remove_subtexture(package_dir, manifest, sub_atlas):
-    filename = manifest['name'].split(os.path.sep)[-1] + '.png'
-    file = package_dir / 'textures/atlases' / filename
-    img = Image.open(file)
+def find_subtextures(pkg_reader):
+    for entry in pkg_reader:
+        if entry.typeName != 'texture':
+            continue
+        
+        for sub in entry.manifest_entry.subAtlases:
+            subname = sub['name']
+            if 'aphrodite' in subname.lower():
+                if subname in ENTRIES:
+                    if ENTRIES[subname]:
+                        yield entry, sub
+                else:
+                    print(f'WARNING: Found new subtexture: {fullpath}. This probably means that Supergiant has added new content to the game. This should probably be re-examined.')
 
+
+def remove_subtexture(path, sub_atlas):
+    img = Image.open(path)
     draw = ImageDraw.Draw(img)
+
     r = sub_atlas['rect']
     top_left = (r['x'], r['y'])
     bottom_right = (r['x'] + r['width'], r['y'] + r['height'])
     draw.rectangle(xy=(top_left, bottom_right), fill=(0, 0, 0, 0))
 
-    img.save(file)
+    img.save(path)
 
 
-def process_pkg(src_path, include, exclude):
-    print('Processing package', src_path,)
+def process_pkg(pkg_path):
     with tempfile.TemporaryDirectory() as tmp:
         tmp = pathlib.Path(tmp)
-        pkg_dir = str(tmp / 'package')
-        print('Extracting package - this may take some time.')
-        deppth.extract(src_path, pkg_dir)
+        extracted = set()
+        print('  Extracting textures')
+        with sggpio.PackageWithManifestReader(pkg_path) as reader:
+            pkg_tmp = tmp / 'package'
+            for entry, sub in find_subtextures(reader):
+                entry.manifest_entry.extract(pkg_tmp)
+                entry.extract(str(pkg_tmp))
 
-        manifests = []
-        for m in (tmp / 'package/manifest').iterdir():
-            with open(m) as f:
-                manifests.append(json.load(f))
+                filename = pkg_tmp / f'textures/atlases/{entry.short_name()}.png'
+                remove_subtexture(filename, sub)
 
-        print('Removing subtextures')
-        modified = []
-        for manifest in manifests:
-            for sub in manifest['subAtlases']:
-                if 'aphrodite' in sub['name'].lower():
-                    if sub['name'] in include:
-                        remove_subtexture(tmp / 'package', manifest, sub)
-                        manifest_name = manifest['name'].split(os.path.sep)[-1]
-                        modified.append(manifest_name)
-                    elif sub['name'] not in exclude:
-                        print(f'WARNING: Unknown texture "{sub["name"]}". You may wish to add this texture to include.txt or exclude.txt.')
-
-        print('Constructing patches')
-        patch_path = str(tmp / 'patch.pkg')
-        deppth.pack(pkg_dir, patch_path, *modified)
-        print('Repacking package')
-        deppth.patch(src_path, patch_path)
-        print('Finished processing package.')
+        print('  Patching package')
+        patch_path = tmp / 'patch.pkg'
+        deppth.pack(str(pkg_tmp), str(patch_path))
+        deppth.patch(pkg_path, patch_path)
 
 
 def run():
@@ -71,12 +70,12 @@ def run():
         print(f'Could not find Hades game files. Is {steam_dir} the correct location for your Steam library?')
         return 1
 
-    include = get_list('include.txt')
-    exclude = get_list('exclude.txt')
     for pkg in 'GUI.pkg', '720p/GUI.pkg', 'BC3/GUI.pkg':
         path = str(base_dir / pkg)
-        process_pkg(path, include, exclude)
+        print('Processing package:', pkg)
+        process_pkg(path)
         print()
+
     print('Finished.')
 
 
diff --git a/setup.py b/setup.py
index 975fc29..7d324eb 100644
--- a/setup.py
+++ b/setup.py
@@ -6,7 +6,8 @@ setup(
     author='Joseph Montanaro',
     install_requires=[
         'deppth@https://github.com/quaerus/deppth/releases/download/v0.1.0.0/deppth-0.1.0.0-py3-none-any.whl',
-        'pillow'
+        'pillow',
+        'lz4',
     ],
     packages=find_packages(),
     entry_points = {