From eb06c114d2f355790cb4c764a42ad4aee18d15b8 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 21 Oct 2025 23:15:48 +0200 Subject: [PATCH] feat: add Snapshot class --- src/snapshot.py | 64 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 src/snapshot.py diff --git a/src/snapshot.py b/src/snapshot.py new file mode 100644 index 0000000..de7bda9 --- /dev/null +++ b/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