From 8c0635bb2a7841d2a361a072bc33f989deb7de44 Mon Sep 17 00:00:00 2001
From: Salad Dais <SaladDais@users.noreply.github.com>
Date: Mon, 18 Jul 2022 06:37:20 +0000
Subject: [PATCH] Add classmethod for rebuilding TEs into a TECollection

---
 hippolyzer/lib/base/templates.py | 29 ++++++++++++++++++++++++++---
 tests/proxy/test_templates.py    |  5 +++++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/hippolyzer/lib/base/templates.py b/hippolyzer/lib/base/templates.py
index 76e7f43..e0416e8 100644
--- a/hippolyzer/lib/base/templates.py
+++ b/hippolyzer/lib/base/templates.py
@@ -3,6 +3,7 @@ Serialization templates for structures used in LLUDP and HTTP bodies.
 """
 
 import abc
+import collections
 import dataclasses
 import enum
 import importlib
@@ -863,7 +864,7 @@ class ShineLevel(IntEnum):
     HIGH = 3
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(unsafe_hash=True)
 class BasicMaterials:
     # Meaning is technically implementation-dependent, these are in LL data files
     Bump: int = se.bitfield_field(bits=5)
@@ -882,7 +883,7 @@ class TexGen(IntEnum):
     CYLINDRICAL = 0x6
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(unsafe_hash=True)
 class MediaFlags:
     WebPage: bool = se.bitfield_field(bits=1, adapter=se.BoolAdapter())
     TexGen: "TexGen" = se.bitfield_field(bits=2, adapter=se.IntEnum(TexGen))
@@ -1129,7 +1130,9 @@ class TextureEntryCollection:
         Makes it easier to get all TE details associated with a specific face
         """
         as_dicts = [dict() for _ in range(num_faces)]
-        for key, vals in dataclasses.asdict(self).items():
+        for field in dataclasses.fields(self):
+            key = field.name
+            vals = getattr(self, key)
             # Fill give all faces the default value for this key
             for te in as_dicts:
                 te[key] = vals[None]
@@ -1144,6 +1147,26 @@ class TextureEntryCollection:
                     as_dicts[face_num][key] = val
         return [TextureEntry(**x) for x in as_dicts]
 
+    @classmethod
+    def from_tes(cls, tes: List[TextureEntry]) -> "TextureEntryCollection":
+        instance = cls()
+        if not tes:
+            return instance
+
+        for field in dataclasses.fields(cls):
+            te_vals: Dict[Any, List[int]] = collections.defaultdict(list)
+            for i, te in enumerate(tes):
+                # Group values by what face they occur on
+                te_vals[getattr(te, field.name)].append(i)
+            # Make most common value the "default", everything else is an exception
+            sorted_vals = sorted(te_vals.items(), key=lambda x: len(x[1]), reverse=True)
+            default_val = sorted_vals.pop(0)[0]
+            te_vals = {None: default_val}
+            for val, face_nums in sorted_vals:
+                te_vals[tuple(face_nums)] = val
+            setattr(instance, field.name, te_vals)
+        return instance
+
 
 TE_SERIALIZER = se.Dataclass(TextureEntryCollection)
 
diff --git a/tests/proxy/test_templates.py b/tests/proxy/test_templates.py
index a0a9b1b..b251090 100644
--- a/tests/proxy/test_templates.py
+++ b/tests/proxy/test_templates.py
@@ -27,6 +27,11 @@ class TemplateTests(unittest.TestCase):
         with self.assertRaises(ValueError):
             deserialized.realize(3)
 
+    def test_tecollection_from_tes(self):
+        deserialized: TextureEntryCollection = TextureEntrySubfieldSerializer.deserialize(None, EXAMPLE_TE)
+        # The TE collection should re-serialize to the same collection when split up and regrouped
+        self.assertEqual(deserialized, TextureEntryCollection.from_tes(deserialized.realize(4)))
+
     def test_face_bitfield_round_trips(self):
         test_val = b"\x81\x03"
         reader = se.BufferReader("!", test_val)