51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
import lzma
|
|
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: lzma.LZMAFile = lzma.LZMAFile(self.path, self.mode)
|
|
|
|
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
|