feat: add record saving
This commit is contained in:
49
src/record_file.py
Normal file
49
src/record_file.py
Normal file
@@ -0,0 +1,49 @@
|
||||
from pathlib import Path
|
||||
import struct
|
||||
import time
|
||||
from typing import Literal
|
||||
|
||||
from src.snapshot import Snapshot
|
||||
|
||||
|
||||
class RecordFile:
|
||||
VERSION = 1
|
||||
|
||||
def __init__(self, path: str | Path, mode: Literal["w", "r"]) -> None:
|
||||
self.path: str | Path = path
|
||||
self.mode: Literal["w", "r"] = mode
|
||||
self.file = open(self.path, self.mode + "b")
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, type, value, traceback):
|
||||
self.file.close()
|
||||
|
||||
def write_header(self, n_snapshots: int):
|
||||
data: bytes = struct.pack(
|
||||
">IId", self.VERSION, n_snapshots, time.time())
|
||||
self.file.write(data)
|
||||
|
||||
def write_snapshots(self, snapshots: list[Snapshot]):
|
||||
self.write_header(len(snapshots))
|
||||
for snapshot in snapshots:
|
||||
data: bytes = snapshot.pack()
|
||||
self.file.write(struct.pack(">I", len(data)) + data)
|
||||
|
||||
def read_snapshots(self) -> list[Snapshot]:
|
||||
version: int = struct.unpack(">I", self.file.read(4))[0]
|
||||
if version != self.VERSION:
|
||||
raise ValueError(
|
||||
f"Cannot parse record file with format version {version} (current version: {self.VERSION})")
|
||||
|
||||
n_snapshots: int
|
||||
timestamp: float
|
||||
n_snapshots, timestamp = struct.unpack(">Id", self.file.read(12))
|
||||
snapshots: list[Snapshot] = []
|
||||
|
||||
for _ in range(n_snapshots):
|
||||
size: int = struct.unpack(">I", self.file.read(4))[0]
|
||||
snapshots.append(Snapshot.unpack(self.file.read(size)))
|
||||
|
||||
return snapshots
|
||||
Reference in New Issue
Block a user