From 0197ba993fcbcb822b654596a97f180658c3651e Mon Sep 17 00:00:00 2001 From: snake-biscuits <36507175+snake-biscuits@users.noreply.github.com> Date: Mon, 8 Jul 2024 03:12:43 +0100 Subject: [PATCH] (utils.binary) replacing old funcs with utils --- bsp_tool/extensions/archives/pi_studios.py | 11 ++------ bsp_tool/extensions/archives/respawn.py | 31 ++++++---------------- bsp_tool/extensions/stbsp.py | 11 +++----- 3 files changed, 13 insertions(+), 40 deletions(-) diff --git a/bsp_tool/extensions/archives/pi_studios.py b/bsp_tool/extensions/archives/pi_studios.py index 808b6aef..370c647a 100644 --- a/bsp_tool/extensions/archives/pi_studios.py +++ b/bsp_tool/extensions/archives/pi_studios.py @@ -2,19 +2,12 @@ from __future__ import annotations import io import struct -from typing import Any, List +from typing import List +from ...utils.extensions.binary import read_struct from . import base -def read_struct(file, format_: str) -> List[Any]: - out = struct.unpack(format_, file.read(struct.calcsize(format_))) - if len(out) == 1: - return out[0] - else: - return out - - class Bpk(base.Archive): filename: str headers: List[CentralHeader] diff --git a/bsp_tool/extensions/archives/respawn.py b/bsp_tool/extensions/archives/respawn.py index 6324a614..eb0d292a 100644 --- a/bsp_tool/extensions/archives/respawn.py +++ b/bsp_tool/extensions/archives/respawn.py @@ -2,27 +2,12 @@ import enum import io import os -import struct -from typing import Any, Dict, List, Tuple, Union +from typing import Dict, List, Tuple +from ...utils.binary import read_str, read_struct from . import base -def read_str(binary_stream: io.BytesIO, encoding="utf-8", errors="strict") -> str: - """for tree parsing""" - out = b"" - c = binary_stream.read(1) - while c != b"\0": - out += c - c = binary_stream.read(1) - return out.decode(encoding, errors) - - -def read_struct(file, format_: str) -> Union[Any, List[Any]]: - out = struct.unpack(format_, file.read(struct.calcsize(format_))) - return out[0] if len(out) == 1 else out - - # TODO: determine actual partial flags (e.g. 0x01 == UI) class RPakTypev7(enum.IntFlag): """observed, could be wrong in some way""" @@ -172,7 +157,7 @@ def __repr__(self): @classmethod def from_stream(cls, stream: io.BytesIO) -> VpkHeader: - magic, major, minor, tree_length, data_length = struct.unpack("I2H2I", stream.read(16)) + magic, major, minor, tree_length, data_length = read_struct(stream, "I2H2I") assert magic == cls.magic, "invalid file magic" assert major == cls.version[0], "unsupported major version" assert minor == cls.version[1], "unsupported minor version" @@ -195,7 +180,7 @@ def __repr__(self): @classmethod def from_stream(cls, vpk_file: io.BytesIO) -> VpkEntry: out = cls() - out.crc, out.preload_length = struct.unpack("IH", vpk_file.read(6)) + out.crc, out.preload_length = read_struct(vpk_file, "IH") # file parts file_part = VpkFilePart.from_stream(vpk_file) while not file_part.is_terminator: @@ -219,12 +204,12 @@ class VpkFilePart: @classmethod def from_stream(cls, vpk_file: io.BytesIO) -> VpkFilePart: out = cls() - out.archive_index = int.from_bytes(vpk_file.read(2), "little") + out.archive_index = read_struct(vpk_file, "H") if out.archive_index == 0xFFFF: out.is_terminator = True return out - out.load_flags = int.from_bytes(vpk_file.read(2), "little") # uint16_t - out.texture_flags = int.from_bytes(vpk_file.read(4), "little") # uint32_t - out.entry_offset, out.entry_length, out.entry_length_uncompressed = struct.unpack("3Q", vpk_file.read(24)) + out.load_flags = read_struct(vpk_file, "H") + out.texture_flags = read_struct(vpk_file, "I") + out.entry_offset, out.entry_length, out.entry_length_uncompressed = read_struct(vpk_file, "3Q") out.is_compressed = (out.entry_length != out.entry_length_uncompressed) return out diff --git a/bsp_tool/extensions/stbsp.py b/bsp_tool/extensions/stbsp.py index e6f72dd6..288fcebf 100644 --- a/bsp_tool/extensions/stbsp.py +++ b/bsp_tool/extensions/stbsp.py @@ -4,15 +4,9 @@ import itertools import os import struct -from typing import Any, Dict, List, Tuple +from typing import Dict, List, Tuple - -def read_struct(file, format_: str) -> List[Any]: - return struct.unpack(format_, file.read(struct.calcsize(format_))) - - -def write_struct(file, format_: str, *args): - file.write(struct.pack(format_, *args)) +from ..utils.binary import read_struct, write_struct class Block(enum.Enum): @@ -145,6 +139,7 @@ def __init__(self, filename: str = "./untitled.stbsp"): self.blocks = {b: BlockIndex(0, 0) for b in Block} # data self.materials = dict() + self.material_infos = list() self.vtfs = list() self.vmts = list() self.columns = list()