Merge pull request #101 from randovania/feature/maya-spline

Properly decode Spline properties

parse_pwe_templates.py

@@ -812,78 +812,40 @@ def game(cls) -> Game:
 
  core_path.joinpath("Color.py").write_text(
  f"""# Generated file
-import dataclasses
 import struct
 import typing
 
 from retro_data_structures.game_check import Game
-from retro_data_structures.properties.base_property import BaseProperty
+from retro_data_structures.properties.base_color import BaseColor
 
 
-@dataclasses.dataclass()
-class Color(BaseProperty):
- r: float = 0.0
- g: float = 0.0
- b: float = 0.0
- a: float = 0.0
-
+class Color(BaseColor):
  @classmethod
  def from_stream(cls, data: typing.BinaryIO, size: typing.Optional[int] = None):
  return cls(*struct.unpack('{endianness}ffff', data.read(16)))
 
  def to_stream(self, data: typing.BinaryIO):
  data.write(struct.pack('{endianness}ffff', self.r, self.g, self.b, self.a))
 
- @classmethod
- def from_json(cls, data: dict):
- return cls(data["r"], data["g"], data["b"], data["a"])
-
- def to_json(self) -> dict:
- return {{
- "r": self.r,
- "g": self.g,
- "b": self.b,
- "a": self.a,
- }}
 """
  + game_code
  )
  core_path.joinpath("Vector.py").write_text(
  f"""# Generated file
-import dataclasses
 import struct
 import typing
 
 from retro_data_structures.game_check import Game
-from retro_data_structures.properties.base_property import BaseProperty
+from retro_data_structures.properties.base_vector import BaseVector
 
 
-@dataclasses.dataclass()
-class Vector(BaseProperty):
- x: float = 0.0
- y: float = 0.0
- z: float = 0.0
-
+class Vector(BaseVector):
  @classmethod
  def from_stream(cls, data: typing.BinaryIO, size: typing.Optional[int] = None):
  return cls(*struct.unpack('{endianness}fff', data.read(12)))
 
  def to_stream(self, data: typing.BinaryIO):
  data.write(struct.pack('{endianness}fff', self.x, self.y, self.z))
-
- @classmethod
- def from_json(cls, data: dict):
- return cls(data["x"], data["y"], data["z"])
-
- def to_json(self) -> dict:
- return {{
- "x": self.x,
- "y": self.y,
- "z": self.z,
- }}
-
- def dependencies_for(self, asset_manager):
- yield from []
 """
  + game_code
  )
@@ -963,7 +925,7 @@ def to_json(self) -> dict:
 
 @dataclasses.dataclass()
 class AnimationParameters(BaseProperty):
- ancs: AssetId = default_asset_id
+ ancs: AssetId = dataclasses.field(metadata={{'asset_types': ['ANCS']}}, default=default_asset_id)
  character_index: int = 0
  initial_anim: int = 0
 
@@ -993,33 +955,69 @@ def dependencies_for(self, asset_manager):
  core_path.joinpath("Spline.py").write_text(
  """# Generated file
 import dataclasses
+import struct
 import typing
-import base64
 
+import construct
+
+from retro_data_structures.common_types import MayaSpline
 from retro_data_structures.game_check import Game
-from retro_data_structures.properties.base_property import BaseProperty
+from retro_data_structures.properties.base_spline import BaseSpline, Knot
+
+
+def _read_knot(data: typing.BinaryIO) -> Knot:
+ header = struct.unpack(">ffBB", data.read(10))
+ cached_tangents_a = None
+ cached_tangents_b = None
+ if header[2] == 5:
+ cached_tangents_a = struct.unpack(">ff", data.read(8))
+ if header[3] == 5:
+ cached_tangents_b = struct.unpack(">ff", data.read(8))
+
+ return Knot(*header, cached_tangents_a=cached_tangents_a, cached_tangents_b=cached_tangents_b)
 
 
 @dataclasses.dataclass()
-class Spline(BaseProperty):
- data: bytes = b""
+class Spline(BaseSpline):
 
  @classmethod
  def from_stream(cls, data: typing.BinaryIO, size: typing.Optional[int] = None):
- assert size is not None
- result = cls()
- result.data = data.read(size)
- return result
+ pre_infinity, post_infinity, knot_count = struct.unpack(">BBL", data.read(6))
+ knots = [
+ _read_knot(data)
+ for _ in range(knot_count)
+ ]
+ clamp_mode, minimum_amplitude, maximum_amplitude = struct.unpack(">Bff", data.read(9))
 
- def to_stream(self, data: typing.BinaryIO):
- data.write(self.data)
+ return cls(
+ pre_infinity=pre_infinity,
+ post_infinity=post_infinity,
+ knots=knots,
+ clamp_mode=clamp_mode,
+ minimum_amplitude=minimum_amplitude,
+ maximum_amplitude=maximum_amplitude,
+ )
 
- @classmethod
- def from_json(cls, data):
- return cls(base64.b64decode(data))
+ def to_stream(self, data: typing.BinaryIO):
+ MayaSpline.build_stream(construct.Container(
+ pre_infinity=self.pre_infinity,
+ post_infinity=self.post_infinity,
+ knots=[
+ construct.Container(
+ time=knot.time,
+ amplitude=knot.amplitude,
+ unk_a=knot.unk_a,
+ unk_b=knot.unk_b,
+ cached_tangents_a=knot.cached_tangents_a,
+ cached_tangents_b=knot.cached_tangents_b,
+ )
+ for knot in self.knots
+ ],
+ clamp_mode=self.clamp_mode,
+ minimum_amplitude=self.minimum_amplitude,
+ maximum_amplitude=self.maximum_amplitude,
+ ), data)
 
- def to_json(self) -> str:
- return base64.b64encode(self.data).decode("ascii")
 """
  + game_code
  )

src/retro_data_structures/common_types.py

@@ -9,7 +9,7 @@
  CString,
  Float32b,
  Hex,
- Int16ub,
+ Int8ub,
  Int32ub,
  Int64ub,
  PaddedString,
@@ -28,6 +28,8 @@
  amplitude=Float32b,
  unk_a=Byte,
  unk_b=Byte,
+ cached_tangents_a=construct.If(construct.this.unk_a == 5, Vector2f),
+ cached_tangents_b=construct.If(construct.this.unk_b == 5, Vector2f),
 )
 
 AABox = Struct(
@@ -67,9 +69,10 @@
 )
 
 MayaSpline = Struct(
- unk=Int16ub,
+ pre_infinity=Int8ub,
+ post_infinity=Int8ub,
  knots=PrefixedArray(Int32ub, Knot),
- clampMode=Byte,
- minAmplitude=Float32b,
- maxAmplitude=Float32b,
+ clamp_mode=Byte,
+ minimum_amplitude=Float32b,
+ maximum_amplitude=Float32b,
 )

src/retro_data_structures/properties/base_color.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import dataclasses
+
+from retro_data_structures.properties.base_property import BaseProperty
+
+
+@dataclasses.dataclass()
+class BaseColor(BaseProperty):
+ r: float = 0.0
+ g: float = 0.0
+ b: float = 0.0
+ a: float = 0.0
+
+ @classmethod
+ def from_json(cls, data: dict):
+ return cls(data["r"], data["g"], data["b"], data["a"])
+
+ def to_json(self) -> dict:
+ return {
+ "r": self.r,
+ "g": self.g,
+ "b": self.b,
+ "a": self.a,
+ }

src/retro_data_structures/properties/base_spline.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import dataclasses
+import typing
+
+from retro_data_structures.properties.base_property import BaseProperty
+
+
+@dataclasses.dataclass()
+class Knot:
+ time: float # X position
+ amplitude: float # Y position
+ unk_a: int
+ unk_b: int
+ cached_tangents_a: tuple[float, float] | None
+ cached_tangents_b: tuple[float, float] | None
+
+ @classmethod
+ def from_json(cls, data: dict) -> typing.Self:
+ return cls(
+ time=data["time"],
+ amplitude=data["amplitude"],
+ unk_a=data["unk_a"],
+ unk_b=data["unk_b"],
+ cached_tangents_a=data["cached_tangents_a"],
+ cached_tangents_b=data["cached_tangents_b"],
+ )
+
+ def to_json(self) -> dict:
+ return {
+ "time": self.time,
+ "amplitude": self.amplitude,
+ "unk_a": self.unk_a,
+ "unk_b": self.unk_b,
+ "cached_tangents_a": self.cached_tangents_a,
+ "cached_tangents_b": self.cached_tangents_b,
+ }
+
+
+@dataclasses.dataclass()
+class BaseSpline(BaseProperty):
+ knots: list[Knot] = dataclasses.field(default_factory=list)
+ minimum_amplitude: float = 0.0
+ maximum_amplitude: float = 0.0
+ pre_infinity: int = 0
+ post_infinity: int = 0
+ clamp_mode: int = 0
+
+ @classmethod
+ def from_json(cls, data: dict) -> typing.Self:
+ return cls(
+ knots=[Knot.from_json(knot) for knot in data["knots"]],
+ minimum_amplitude=data["minimum_amplitude"],
+ maximum_amplitude=data["maximum_amplitude"],
+ pre_infinity=data["pre_infinity"],
+ post_infinity=data["post_infinity"],
+ clamp_mode=data["clamp_mode"],
+ )
+
+ def to_json(self) -> dict:
+ return {
+ "knots": [knot.to_json() for knot in self.knots],
+ "minimum_amplitude": self.minimum_amplitude,
+ "maximum_amplitude": self.maximum_amplitude,
+ "pre_infinity": self.pre_infinity,
+ "post_infinity": self.post_infinity,
+ "clamp_mode": self.clamp_mode,
+ }

src/retro_data_structures/properties/base_vector.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import dataclasses
+
+from retro_data_structures.properties.base_property import BaseProperty
+
+
+@dataclasses.dataclass()
+class BaseVector(BaseProperty):
+ x: float = 0.0
+ y: float = 0.0
+ z: float = 0.0
+
+ @classmethod
+ def from_json(cls, data: dict):
+ return cls(data["x"], data["y"], data["z"])
+
+ def to_json(self) -> dict:
+ return {
+ "x": self.x,
+ "y": self.y,
+ "z": self.z,
+ }
+
+ def dependencies_for(self, asset_manager):
+ yield from []