From 6c951819707658d1663bc5f5a23f7cefe9d4a529 Mon Sep 17 00:00:00 2001
From: dyceron <dycerontv@gmail.com>
Date: Sun, 22 Dec 2024 16:33:15 -0500
Subject: [PATCH] Address reviews

---
 .../formats/bmtun.py                          | 26 ++++++++++---------
 tests/formats/test_bmtun.py                   |  6 ++---
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/src/mercury_engine_data_structures/formats/bmtun.py b/src/mercury_engine_data_structures/formats/bmtun.py
index 3652bf2..5d60168 100644
--- a/src/mercury_engine_data_structures/formats/bmtun.py
+++ b/src/mercury_engine_data_structures/formats/bmtun.py
@@ -1,13 +1,12 @@
 from __future__ import annotations
 
 import functools
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import construct
 from construct.core import (
     Const,
     Construct,
-    Container,
     Flag,
     Int32sl,
     Struct,
@@ -50,14 +49,17 @@ class Bmtun(BaseResource):
     def construct_class(cls, target_game: Game) -> Construct:
         return BMTUN
 
-    def get_tunable(self, tunable: str, param: str) -> Container:
+    def _check_tunable_exists(self, class_name: str, tunable_name: str) -> None:
         classes = self.raw.classes
-        if tunable not in classes:
-            raise ValueError(f"Unknown tunable: {tunable}!")
-        if param not in classes[tunable].tunables:
-            raise ValueError(f"Unknown tunable param: {param}!")
-        return self.raw.classes[tunable].tunables[param]
-
-    def set_tunable(self, tunable_name: str, param: str, value: Any) -> None:
-        tunable = self.get_tunable(tunable_name, param)
-        tunable.value = value
+        if class_name not in classes:
+            raise KeyError(f"Unknown tunable: {class_name}!")
+        if tunable_name not in classes[class_name].tunables:
+            raise KeyError(f"Unknown tunable param: {tunable_name}!")
+
+    def get_tunable(self, class_name: str, tunable_name: str) -> None:
+        self._check_tunable_exists(class_name, tunable_name)
+        return self.raw.classes[class_name].tunables[tunable_name].value
+
+    def set_tunable(self, class_name: str, tunable_name: str, value: str | float | bool | int | list[float]) -> None:
+        self._check_tunable_exists(class_name, tunable_name)
+        self.raw.classes[class_name].tunables[tunable_name].value = value
diff --git a/tests/formats/test_bmtun.py b/tests/formats/test_bmtun.py
index d318f90..1275399 100644
--- a/tests/formats/test_bmtun.py
+++ b/tests/formats/test_bmtun.py
@@ -18,12 +18,12 @@ def bmtun(samus_returns_tree) -> Bmtun:
 
 
 def test_get_tunable(bmtun):
-    assert bmtun.get_tunable("Amiibo|CTunableReserveTanks", "fLifeTankSize").value == 299.0
+    assert bmtun.get_tunable("Amiibo|CTunableReserveTanks", "fLifeTankSize") == 299.0
 
 
 def test_set_tunable(bmtun):
     bmtun.set_tunable("Amiibo|CTunableReserveTanks", "fLifeTankSize", 199.0)
     assert bmtun.get_tunable("Amiibo|CTunableReserveTanks", "fLifeTankSize") == 199.0
 
-    with pytest.raises(ValueError):
-        bmtun.set_tunable("Amiibo|CTunableReserveTanks", "fSETankSize", True)
+    bmtun.set_tunable("CTunableMissile", "sDamageSource", "BOMB")
+    assert bmtun.get_tunable("CTunableMissile", "sDamageSource") == "BOMB"