feat: add Snapshot class

src/snapshot.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import struct
+from dataclasses import dataclass, field
+from typing import Optional
+
+import numpy as np
+
+from vec import Vec
+
+
+def iter_unpack(format, data):
+    nbr_bytes = struct.calcsize(format)
+    return struct.unpack(format, data[:nbr_bytes]), data[nbr_bytes:]
+
+
+@dataclass
+class Snapshot:
+    controls: tuple[bool, bool, bool, bool] = (False, False, False, False)
+    position: Vec = field(default_factory=Vec)
+    direction: Vec = field(default_factory=Vec)
+    speed: float = 0
+    raycast_distances: list[float] | tuple[float, ...] = [0]
+    image: Optional[np.ndarray] = None
+
+    def pack(self):
+        data: bytes = b""
+        data += struct.pack(">BBBB", *self.controls)
+        data += struct.pack(
+            ">fffff",
+            self.position.x,
+            self.position.y,
+            self.direction.x,
+            self.direction.y,
+            self.speed,
+        )
+
+        nbr_raycasts: int = len(self.raycast_distances)
+        data += struct.pack(f">B{nbr_raycasts}f", nbr_raycasts, *self.raycast_distances)
+
+        if self.image is not None:
+            data += struct.pack(">II", self.image.shape[0], self.image.shape[1])
+            data += self.image.tobytes()
+        else:
+            data += struct.pack(">II", 0, 0)
+
+        return data
+
+    def unpack(self, data):
+        self.controls, data = iter_unpack(">BBBB", data)
+        (x, y, dx, dy, s), data = iter_unpack(">fffff", data)
+        self.position = Vec(x, y)
+        self.direction = Vec(dx, dy)
+        self.speed = s
+
+        (nbr_raycasts,), data = iter_unpack(">B", data)
+        self.raycast_distances, data = iter_unpack(f">{nbr_raycasts}f", data)
+
+        (h, w), data = iter_unpack(">ii", data)
+
+        if h * w > 0:
+            self.image = np.frombuffer(data, np.uint8).reshape(h, w, 3)
+        else:
+            self.image = None