diff --git a/.github/workflows/unittest.yaml b/.github/workflows/unittest.yaml index 556c570..4af3352 100644 --- a/.github/workflows/unittest.yaml +++ b/.github/workflows/unittest.yaml @@ -8,7 +8,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.8", "3.9", "3.10", "3.11"] + python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] steps: - uses: actions/checkout@v3 diff --git a/.gitignore b/.gitignore index e3ad8b9..53c3d5e 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,5 @@ dist/ env/ pyBIG.egg-info/ tests/test_data/output/* +test.big +test.py diff --git a/LICENSE b/LICENSE index 63b4b68..2f90256 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) [year] [fullname] +Copyright (c) 2025 Clement Julia Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index c9548e1..02c3e36 100644 --- a/README.md +++ b/README.md @@ -12,24 +12,27 @@ pip install pyBIG ``` ## Usage -The library is based on the pyBIG.Archive object. This objects takes raw bytes representing a BIG archive. The decision to take raw bytes allow the user to decide where those bytes come from, whether a file stored in memory or on disk. There is also a class method, Archive.from_directory that allows you to load a directory on the disk painlessly. +This library offers a few different implementations of BaseArchive that all represent a .BIG archive. Their main difference is how they manipulate the data. Read below to select the best one for your use case. All these objects have the same or very similar interface. Namely: + - BaseArchive.edit_file(str, bytes) + - BaseArchive.add_file(str, bytes) + - BaseArchive.remove_file(str) -You can modify the archive in memory with the following methods: - - Archive.edit_file(str, bytes) - - Archive.add_file(str, bytes) - - Archive.remove_file(str) +Each method takes a name which is the windows-format path to the file in the archive so something like 'data\ini\weapon.ini'. The methods that takes bytes represent the new contents of the file as bytes. To apply the changes you need to use BaseARchuve.repack(). -Each method takes a name which is the windows-format path to the file in the archive so something like 'data\ini\weapon.ini'. The methods that takes bytes represent the new contents of the file as bytes. +There are also a few utility functions + - BaseArchive.from_directory(str, str, **kwargs) + - BaseArchive.empty(str, **kwargs) -It is important to note that these methods do not actually modify the archive but it is as if. This does not update the entries or the raw bytes. If you want to update the archive you need to call Archive.repack(). This is an expensive operation which is only called automatically when the archive is saved or extracted. The rest is up to the user. +Below is a more in depth explaination. You can look at the tests for more examples. -You can look at the tests for more examples. +### InMemoryArchive +As the name implies, the InMemoryArchive loads the entire archive into memory and keeps it there, doing all manipulations from there. You can save it back to disk with InMemoryArchive.save(str). ```python -from pyBIG import Archive +from pyBIG import InMemoryArchive with open("test.big", "rb") as f: - archive = Archive(f.read()) + archive = InMemoryArchive(f.read()) # get the contents of a file as bytes contents = archive.read_file("data\\ini\\weapon.ini") @@ -50,20 +53,61 @@ archive.save("test.big") archive.extract("output/") # load an archive from a directory -archive = Archive.from_directory("output/") +archive = InMemoryArchive.from_directory("output/") ``` -### Advanced -In version 0.2.0, this library also makes the `LargeArchive` object available. This special object does not store the entire file into memory, allowing for manipulation of large files. It works essentially the same except that reading is done from the file present on disk and functions are tied to that location. Repacking does the same as save on this object but it is recommended to instead use the save function. +### InDiskArchive +The InDiskArchive does not store the entire file into memory, allowing for manipulation of larger files. It works essentially the same except that reading is done from the file present on disk and functions are tied to that location. Repacking does the same as save on this object but it is recommended to instead use the save function. -It is important to note that adding and editing files in a LargeArchive stores them in memory. As such it is recommended to to save at regular interval to commit these changes to disk. The LargeArchive object exposes `archive_memory_size` as a simple way of seeing how many bytes are currently stored directly on the object. +It is important to note that adding and editing files in a InDiskArchive stores them in memory. As such it is recommended to save at regular interval to commit these changes to disk. The BaseArchive object exposes `archive_memory_size` as a simple way of seeing how many bytes are currently stored directly on the object. ```python -from pyBIG import LargeArchive +from pyBIG import InDiskArchive -archive = LargeArchive("test.big") +archive = InDiskArchive("test.big") ``` +## RefPack + +The library grossly implements the refpack compression algorithm which allows users to compress and decompress files to and from that format. This is done very simply: +```python + +from pyBIG import refpack + +to_compress = b"My bytes to compress" +compressed = refpack.compress(to_compress) +decompressed = refpack.decompress(compressed) + + +assert to_compress == decompressed +``` + +You can also check if data has the refpack header which is a potential indicator that the data is refpack encoded using `refpack.has_refpack_header`. Data without the header could still be encoded, just without the header. Best way to try is to just attempt to decompress, python zen and all. + +For clarity, you must compressed individual files before adding them to the the .big file, is is entirely left up to the reponsibility of the user to do this. If you have done so then the SAGE engine games will be able to read the compressed files flawlessly. + +## Tests + +Tests must be run from root directory +* `python -m unittest tests.functional_tests` +* `python -m unittest tests.memory_tests` +* `python -m unittest tests.profiler` + + ## TODO -- [ ] Investigate and implement proper compression (refpack) +- [x] Investigate and implement proper compression (refpack) + + +## Changelog + +### v0.6.0 +- Archive renamed to InMemoryArchive (alias remains for backwards compatibility) +- LargeArchive renamed to InDiskArchive (alias remains for backward comaptibility) +- Backend reworked to be cleaner +- Archives now handle different .big types +- `InDiskArchive.from_directory` implemented but not very efficient yet +- Added more typing +- Added `BaseArchive.bytes` +- Inmplemented refpack compression + diff --git a/pyBIG/__init__.py b/pyBIG/__init__.py index 73dafcb..723124f 100644 --- a/pyBIG/__init__.py +++ b/pyBIG/__init__.py @@ -1,6 +1,9 @@ -from .archive import Archive -from .large_archive import LargeArchive +from .memory_archive import InMemoryArchive +from .disk_archive import InDiskArchive -__version__ = "0.5.0" +Archive = InMemoryArchive +LargeArchive = InDiskArchive -__all__ = ["Archive", "LargeArchive"] +__version__ = "0.6.0" + +__all__ = ["InMemoryArchive", "InDiskArchive", "Archive", "LargeArchive"] diff --git a/pyBIG/archive.py b/pyBIG/archive.py deleted file mode 100644 index feb708c..0000000 --- a/pyBIG/archive.py +++ /dev/null @@ -1,212 +0,0 @@ -import io -import logging -import os -import struct - -from .base_archive import BaseArchive, Entry, FileAction - - -class Archive(BaseArchive): - """The core of the library, represents a BIG file and allows - the user to mainpulate it programatically - - Params - ------- - content : Optional[bytes] - Raw bytes of the original big file - - """ - - def __init__(self, content: bytes = b"", *, entries=None): - self.archive = io.BytesIO(content) - - self.header = self.archive.read(4).decode("utf-8") - - self.entries = entries or {} - self.modified_entries = {} - - if entries is None: - self.entries = self._unpack(self.archive) - - def __repr__(self): - return f"< Archive entries={len(self.entries)} dirty={bool(self.modified_entries)}" - - def _pack(self): - """Rewrite the archive with the modifications stores - in self.modified_entries.""" - binary, total_size, file_count = self._create_file_list() - self.archive, self.entries = self._pack_file_list(binary, total_size, file_count) - self.archive.seek(0) - self.modified_entries = {} - - @staticmethod - def _pack_file_list(binary, total_size, file_count): - """Index the files and append the raw data to create a complete archive""" - archive = io.BytesIO() - raw_data = io.BytesIO() - entries = {} - - # header, charstring, 4 bytes - always BIG4 or something similiar - header = "BIG4" - archive.write(struct.pack("4s", header.encode("utf-8"))) - - # https://github.com/chipgw/openbfme/blob/master/bigreader/bigarchive.cpp - # /* 8 bytes for every entry + 20 at the start and end. */ - first_entry = (len(binary) * 8) + 20 - - for file in binary: - first_entry += len(file[0]) + 1 - - # total file size, unsigned integer, 4 bytes, little endian byte order - size = total_size + first_entry + 1 - logging.info(f"size: {size}") - archive.write(struct.pack("I", file_count)) - - # total size of index table in bytes, unsigned integer, 4 bytes, big endian byte order - logging.info(f"index size: {first_entry}") - archive.write(struct.pack(">I", first_entry)) - - position = 1 - - logging.info("packing files...") - for file in binary: - # position of embedded file within BIG-file, unsigned integer, 4 bytes, big endian byte order - # size of embedded data, unsigned integer, 4 bytes, big endian byte order - pos_size = struct.pack(">II", first_entry + position, file[1]) - - # file name, cstring, ends with null byte - name = file[0].encode("latin-1") + b"\x00" - packed_name = struct.pack(f"{len(name)}s", name) - archive.write(pos_size + packed_name) - - raw_data.write(file[2]) - - entries[file[0]] = Entry(file[0], first_entry + position, file[1]) - - position += file[1] - - # not sure what's this but I think we need it see: - # https://github.com/chipgw/openbfme/blob/master/bigreader/bigarchive.cpp - archive.write(b"L253") - archive.write(b"\0") - - # raw file data at the positions specified in the index - raw_data.seek(0) - archive.write(raw_data.read()) - logging.debug("DONE") - - return archive, entries - - @staticmethod - def _create_file_list_from_directory(path): - """Gather the necessary information on each file in the directory to - prepare for packing.""" - binary_files = [] - file_count = 0 - total_size = 0 - - for dir_name, _, file_list in os.walk(path): - for filename in file_list: - file_path = os.path.join(dir_name, filename) - name = file_path.replace(path, "")[1:].replace("/", "\\") - - with open(file_path, "rb") as f: - contents = f.read() - size = len(contents) - - logging.debug(f"name: {name}") - logging.debug("position: ???") - logging.debug(f"file size: {size}") - binary_files.append((name, size, contents)) - - file_count += 1 - total_size += size - - binary_files.sort(key=lambda x: x[0]) - return binary_files, total_size, file_count - - def _create_file_list(self): - """Re-gather the necessary information on each file in the archive - while taking into account the modifications made by the user since - then. - """ - binary_files = [] - file_count = 0 - total_size = 0 - - for name in self.entries: - if name in self.modified_entries: - entry = self.modified_entries[name] - entry_bytes = entry.content - - if entry.action is FileAction.REMOVE: - logging.info(f"removing {name}") - continue - - logging.info(f"editing {name}") - else: - entry = self.entries[name] - entry_bytes = self._get_file(name) - - binary_files.append((entry.name, len(entry_bytes), entry_bytes)) - file_count += 1 - total_size += len(entry_bytes) - - for entry in [x for x in self.modified_entries.values() if x.action is FileAction.ADD]: - logging.info(f"adding {entry.name}") - entry_bytes = entry.content - binary_files.append((entry.name, len(entry_bytes), entry_bytes)) - file_count += 1 - total_size += len(entry_bytes) - - binary_files.sort(key=lambda x: x[0]) - return binary_files, total_size, file_count - - def _get_file(self, name: str) -> bytes: - """Get the contents of a specific file in the big based on file name""" - entry = self.entries[name] - self.archive.seek(entry.position) - return self.archive.read(entry.size) - - def save(self, path: str): - """Save the archive to a file. - - Params - ------- - path : str - The path to save to. Something like 'path/to/file/test.big' - """ - self._pack() - with open(path, "wb") as f: - f.write(self.archive.getbuffer()) - - @classmethod - def from_directory(cls, path: str) -> "Archive": - """Generate a BIG archive from a directory. This is useful for - compiling an archive without adding each file manually. You simply - give the top level directory and every file will be added recursively. - - Params - ------- - path : str - Path to the top level folder of the files you wish to compile - - Returns - -------- - Archive - Compiled archived - """ - binary_files, total_size, file_count = cls._create_file_list_from_directory(path) - archive, entries = cls._pack_file_list(binary_files, total_size, file_count) - archive.seek(0) - - return cls(archive.read(), entries=entries) - - @classmethod - def empty(cls): - """Generate an empty archive.""" - return cls(entries={}) diff --git a/pyBIG/base_archive.py b/pyBIG/base_archive.py index f4f9141..c5b9211 100644 --- a/pyBIG/base_archive.py +++ b/pyBIG/base_archive.py @@ -3,10 +3,14 @@ import os import struct from collections import namedtuple -from typing import Dict, List +from typing import IO, Dict, List, Tuple, Type, TypeVar, Union + +from .utils import MaxSizeError Entry = namedtuple("Entry", "name position size") EntryEdit = namedtuple("EntryEdit", "name action content") +FileList = Union[List[Tuple[str, int]], List[Tuple[str, int, int]]] +T = TypeVar("T", bound="BaseArchive") class FileAction(enum.Enum): @@ -17,21 +21,16 @@ class FileAction(enum.Enum): class BaseArchive: modified_entries: Dict[str, EntryEdit] - - def _pack(self): - """Rewrite the archive with the modifications stores - in self.modified_entries.""" - - raise NotImplementedError + entries: Dict[str, Entry] @staticmethod - def _unpack(file): + def _unpack(file: IO) -> Tuple[List[Entry], str]: """Get a list of files in the big""" entries = {} file.seek(0) # header - file.read(4) + header = file.read(4).decode("utf-8") file_size = struct.unpack("I", file.read(4))[0] logging.info(f"size: {file_size}") @@ -57,8 +56,125 @@ def _unpack(file): entries[name] = Entry(name, position, entry_size) + return entries, header + + def _create_file_list(self) -> Tuple[FileList, int, int]: + """Re-gather the necessary information on each file in the archive + while taking into account the modifications made by the user since + then. + """ + file_list = [] + file_count = 0 + total_size = 0 + + for name in self.entries: + if name in self.modified_entries: + entry = self.modified_entries[name] + entry_bytes = entry.content + + if entry.action is FileAction.REMOVE: + logging.info(f"removing {name}") + continue + + logging.info(f"editing {name}") + else: + entry = self.entries[name] + entry_bytes = self._get_file(name) + + file_list.append((entry.name, len(entry_bytes), entry_bytes)) + file_count += 1 + total_size += len(entry_bytes) + + for entry in [x for x in self.modified_entries.values() if x.action is FileAction.ADD]: + logging.info(f"adding {entry.name}") + entry_bytes = entry.content + file_list.append((entry.name, len(entry_bytes), entry_bytes)) + file_count += 1 + total_size += len(entry_bytes) + + file_list.sort(key=lambda x: x[0]) + return file_list, total_size, file_count + + def _pack_file_list( + self, + archive_file: IO, + file_list: List[Tuple[str, int]], + total_size: int, + file_count: int, + header: str, + ): + """Index the files and append the raw data to create a complete archive""" + entries = {} + + # header, charstring, 4 bytes - always BIG4 or something similiar + archive_file.write(struct.pack("4s", header.encode("utf-8"))) + + # https://github.com/chipgw/openbfme/blob/master/bigreader/bigarchive.cpp + # /* 8 bytes for every entry + 20 at the start and end. */ + first_entry = (len(file_list) * 8) + 20 + + for file in file_list: + first_entry += len(file[0]) + 1 + + # total file size, unsigned integer, 4 bytes, little endian byte order + size = total_size + first_entry + 1 + logging.info(f"size: {size}") + + try: + archive_file.write(struct.pack("I", file_count)) + + # total size of index table in bytes, unsigned integer, 4 bytes, big endian byte order + logging.info(f"index size: {first_entry}") + archive_file.write(struct.pack(">I", first_entry)) + + position = 1 + + logging.info("packing file list...") + for file in file_list: + # position of embedded file within BIG-file, unsigned integer, 4 bytes, big endian byte order + # size of embedded data, unsigned integer, 4 bytes, big endian byte order + logging.debug("packing %s", file[0]) + pos_size = struct.pack(">II", first_entry + position, file[1]) + + # file name, cstring, ends with null byte + name = file[0].encode("latin-1") + b"\x00" + packed_name = struct.pack(f"{len(name)}s", name) + archive_file.write(pos_size + packed_name) + + entries[file[0]] = Entry(file[0], first_entry + position, file[1]) + + position += file[1] + + # not sure what's this but I think we need it see: + # https://github.com/chipgw/openbfme/blob/master/bigreader/bigarchive.cpp + archive_file.write(b"L253") + archive_file.write(b"\0") + logging.info("DONE packing file list") + return entries + @staticmethod + def _pack_archive_from_directory(archive: T, path: str) -> T: + logging.info("building archive from folder") + for dir_name, _, file_list in os.walk(path): + for filename in file_list: + file_path = os.path.join(dir_name, filename) + name = file_path.replace(path, "")[1:].replace("/", "\\") + + with open(file_path, "rb") as f: + logging.debug("adding %s", name) + archive.add_file(name, f.read()) + + archive._pack() + logging.info("done building archive from folder") + return archive + def file_exists(self, name: str) -> bool: """Check if a file exists @@ -78,13 +194,14 @@ def file_exists(self, name: str) -> bool: return name in self.entries def file_list(self) -> List[str]: - """Return a list of files, compiling both actual files - and new/removed files + """Return a list of file names, compiling both actual files + currently in the archive and new/removed files waiting + to be repacked. Returns -------- List[str] - The list of files + The list of file names """ file_list = list( { @@ -105,9 +222,9 @@ def file_list(self) -> List[str]: return file_list - def read_file(self, name: str): - """Get the raw bytes of the file if the file exists. This method has the - advantage over simply accessing Archive.entries that it will + def read_file(self, name: str) -> bytes: + """Get the raw bytes of the file if the file exists. This method has + the advantage over simply accessing Archive.entries that it will also check pending modified entries Params @@ -203,7 +320,7 @@ def remove_file(self, name: str): self.modified_entries[name] = EntryEdit(name, FileAction.REMOVE, None) - def extract(self, output: str, *, files: List[str] = ()): + def extract(self, output: str, *, files: List[str] = None): """Extract the contents of the archive to a folder. Params @@ -213,7 +330,7 @@ def extract(self, output: str, *, files: List[str] = ()): files : Optional[List[str]] The list of files to extract """ - if not files: + if files is None: files = self.file_list() for name in files: @@ -232,11 +349,47 @@ def repack(self): """Update the archive to include all the modified entries. This clears the list and updates the archive with the new data. """ - if not self.modified_entries: - return - self._pack() + def archive_memory_size(self) -> int: + """Get the current in memory size of all the modifies entries that + have not yet been saved. You can use this to decide when you would + like to save in relation to the capacities of your machine. + """ + return sum( + [ + len(entry.content) + for entry in self.modified_entries.values() + if entry.action in (FileAction.ADD, FileAction.EDIT) + ] + ) + + def _create_entry(self, name: str, content: bytes) -> tuple: + """Create a file entry.""" + + raise NotImplementedError + + def _get_file(self, name: str) -> bytes: + """Archive specific method for retrieving file bytes from + the archive. + """ + + raise NotImplementedError + + def _pack(self): + """Rewrite the archive with the modifications stored + in self.modified_entries. + """ + + raise NotImplementedError + + def _pack_files( + self, raw_data_file: IO, file_list: FileList, total_size: int, file_count: int + ): + """Combine all files into a single raw data bundle""" + + raise NotImplementedError + def save(self, path: str): """Save the archive to a file. @@ -248,7 +401,7 @@ def save(self, path: str): raise NotImplementedError @classmethod - def from_directory(cls, path: str) -> "BaseArchive": + def from_directory(cls: Type[T], path: str, header: str = "BIG4", **kwargs) -> T: """Generate a BIG archive from a directory. This is useful for compiling an archive without adding each file manually. You simply give the top level directory and every file will be added recursively. @@ -257,6 +410,8 @@ def from_directory(cls, path: str) -> "BaseArchive": ------- path : str Path to the top level folder of the files you wish to compile + header : str + The type of the archive, either BIG4 or BIGF Returns -------- @@ -264,3 +419,31 @@ def from_directory(cls, path: str) -> "BaseArchive": Compiled archived """ raise NotImplementedError + + @classmethod + def empty(cls: Type[T], header: str = "BIG4", **kwargs) -> T: + """Generate an empty archive. + + Params + ------- + header : str + The type of the archive, can either be BIG4 or BIGF. Defaults to BIG4 + + Returns + -------- + Archive + Empty archive + """ + + raise NotImplementedError + + def bytes(self): + """Returns the archive data as bytes + + Returns + -------- + bytes + The archive data + """ + + raise NotImplementedError diff --git a/pyBIG/disk_archive.py b/pyBIG/disk_archive.py new file mode 100644 index 0000000..5bb20af --- /dev/null +++ b/pyBIG/disk_archive.py @@ -0,0 +1,164 @@ +import logging +import os +import shutil +import tempfile +from typing import IO, Type, TypeVar + +from .base_archive import BaseArchive, FileAction, FileList + +T = TypeVar("T", bound="InDiskArchive") + + +class InDiskArchive(BaseArchive): + """This implementation stores as few things possible in memory, preferring + to use disk space in the temp folder instead. This allows loading large + archives into memory with minimal impact on memory usage. Assume all changes made + are directly applied to the data in the disk. + + Params + ------- + file_path : str + The path to the archive. + """ + + def __init__(self, file_path: str, *, entries=None, header: str = "BIG4"): + self.file_path = file_path + self.modified_entries = {} + + if not os.path.exists(file_path): + raise ValueError(f"File {file_path} not found") + + if entries is None: + with open(self.file_path, "rb") as f: + self.entries, self.header = self._unpack(f) + else: + self.entries = entries + self.header = header + + def __repr__(self): + return f"< LargeArchive path={self.file_path} entries={len(self.entries)} dirty={bool(self.modified_entries)} >" + + def _pack(self, file_path=None): + """Rewrite the archive with the modifications stores + in self.modified_entries.""" + file_data = self._create_file_list() + + with tempfile.NamedTemporaryFile(delete=False) as fp: + self.entries = self._pack_file_list(fp, *file_data, self.header) + name = fp.name + + self._pack_files(fp, *file_data) + + path = file_path or self.file_path + shutil.move(name, path) + self.modified_entries = {} + + def _create_entry(self, name: str, content: bytes) -> tuple: + """In this archive we try to keep as few things in memory + as possible, as such an entry does not include the contents + itself. + """ + + return name, len(content) + + def _pack_files( + self, raw_data_file: IO, file_list: FileList, total_size: int, file_count: int + ): + """Combine all files into a single raw data bundle""" + logging.info("packing files") + + # raw file data at the positions specified in the index + with open(self.file_path, "rb") as existing_archive: + for file in file_list: + if file[0] in self.modified_entries: + file_entry = self.modified_entries[file[0]] + if file_entry.action is not FileAction.REMOVE: + raw_data_file.write(file_entry.content) + else: + file_entry = self.entries[file[0]] + existing_archive.seek(file_entry.position) + raw_data_file.write(existing_archive.read(file_entry.size)) + + logging.info("finished packing files") + + def _get_file(self, name: str) -> bytes: + """Get the contents of a specific file in the big based on file name""" + entry = self.entries[name] + with open(self.file_path, "rb") as f: + f.seek(entry.position) + return f.read(entry.size) + + def save(self, path: str = None): + """Save the archive to a file. + + Params + ------- + path : Optional[str] + The new path to save to. Something like 'path/to/file/test.big'. + Omit this if you just want to save in the same file. + """ + self._pack(path) + + @classmethod + def from_directory( + cls: Type[T], path: str, header: str = "BIG4", *, file_path: str = None + ) -> T: + """Generate a BIG archive from a directory. This is useful for + compiling an archive without adding each file manually. You simply + give the top level directory and every file will be added recursively. + + Params + ------- + path : str + Path to the top level folder of the files you wish to compile + header : str + The type of the archive, either BIG4 or BIGF + file_path : str + Path to save the new archive + + Returns + -------- + Archive + Compiled archived + """ + if file_path is None: + raise ValueError("Please specify a file path") + + return cls._pack_archive_from_directory(cls.empty(header, file_path=file_path), path) + + @classmethod + def empty(cls: Type[T], header: str = "BIG4", *, file_path: str = None) -> T: + """Generate an empty archive. + + Params + ------- + header : str + The type of the archive, can either be BIG4 or BIGF. Defaults to BIG4 + file_path : str + Path to save the new archive + + Returns + -------- + Archive + Empty archive + """ + if os.path.exists(file_path): + raise ValueError(f"File {file_path} already exists.") + + with open(file_path, "wb") as f: + f.write(b"") + + return cls(file_path, entries={}, header=header) + + def bytes(self): + """Returns the archive data as bytes + + Returns + -------- + bytes + The archive data + """ + self._pack() + + with open(self.file_path, "rb") as f: + return f.read() diff --git a/pyBIG/large_archive.py b/pyBIG/large_archive.py deleted file mode 100644 index 8951a66..0000000 --- a/pyBIG/large_archive.py +++ /dev/null @@ -1,267 +0,0 @@ -import logging -import os -import shutil -import struct -import tempfile - -from .base_archive import BaseArchive, Entry, FileAction -from .utils import MaxSizeError - - -class LargeArchive(BaseArchive): - """If an archive is too large to load into memory, you can instead use this class. It expects a file path - rather than raw bytes and doesn't store the entire file into memory. Rather it simply reads the headers - to obtain the list file and then stream data from the file when it needs to. You cannot repack a large - archive because repacking implies that no changes are made to the file. Instead use save. - - TODO: Additionally, this class can be used to construct and archive from a large directory. - - Params - ------- - file_path : str - The path to the archive. - """ - - def __init__(self, file_path: str, *, entries=None): - self.file_path = file_path - self.modified_entries = {} - - if not os.path.exists(file_path): - raise ValueError(f"File {file_path} not found") - - if entries is None: - with open(self.file_path, "rb") as f: - self.entries = self._unpack(f) - else: - self.entries = entries - - def __repr__(self): - return f"< LargeArchive path={self.file_path} entries={len(self.entries)} dirty={bool(self.modified_entries)}" - - def _pack(self, file_path=None): - """Rewrite the archive with the modifications stores - in self.modified_entries.""" - binary, total_size, file_count = self._create_file_list() - - with tempfile.NamedTemporaryFile(delete=False) as fp: - self.entries = self._pack_file_list(fp, binary, total_size, file_count) - name = fp.name - - path = file_path or self.file_path - shutil.move(name, path) - self.modified_entries = {} - - def _pack_file_list(self, archive, binary, total_size, file_count): - """Index the files and append the raw data to create a complete archive""" - entries = {} - - # header, charstring, 4 bytes - always BIG4 or something similiar - header = "BIG4" - archive.write(struct.pack("4s", header.encode("utf-8"))) - - # https://github.com/chipgw/openbfme/blob/master/bigreader/bigarchive.cpp - # /* 8 bytes for every entry + 20 at the start and end. */ - first_entry = (len(binary) * 8) + 20 - - for file in binary: - first_entry += len(file[0]) + 1 - - # total file size, unsigned integer, 4 bytes, little endian byte order - size = total_size + first_entry + 1 - logging.info(f"size: {size}") - - try: - archive.write(struct.pack("I", file_count)) - - # total size of index table in bytes, unsigned integer, 4 bytes, big endian byte order - logging.info(f"index size: {first_entry}") - archive.write(struct.pack(">I", first_entry)) - - position = 1 - - logging.info("packing files...") - for file in binary: - # position of embedded file within BIG-file, unsigned integer, 4 bytes, big endian byte order - # size of embedded data, unsigned integer, 4 bytes, big endian byte order - pos_size = struct.pack(">II", first_entry + position, file[1]) - - # file name, cstring, ends with null byte - name = file[0].encode("latin-1") + b"\x00" - packed_name = struct.pack(f"{len(name)}s", name) - archive.write(pos_size + packed_name) - entries[file[0]] = Entry(file[0], first_entry + position, file[1]) - - position += file[1] - - # not sure what's this but I think we need it see: - # https://github.com/chipgw/openbfme/blob/master/bigreader/bigarchive.cpp - archive.write(b"L253") - archive.write(b"\0") - - # raw file data at the positions specified in the index - with open(self.file_path, "rb") as existing_archive: - for file in binary: - if file[0] in self.modified_entries: - file_entry = self.modified_entries[file[0]] - if file_entry.action is not FileAction.REMOVE: - archive.write(file_entry.content) - else: - file_entry = self.entries[file[0]] - existing_archive.seek(file_entry.position) - archive.write(existing_archive.read(file_entry.size)) - - logging.debug("DONE") - - return entries - - @staticmethod - def _create_file_list_from_directory(path): - # TODO: Fix this - """Gather the necessary information on each file in the directory to - prepare for packing.""" - binary_files = [] - file_count = 0 - total_size = 0 - - for dir_name, _, file_list in os.walk(path): - for filename in file_list: - file_path = os.path.join(dir_name, filename) - name = file_path.replace(path, "")[1:].replace("/", "\\") - - with open(file_path, "rb") as f: - contents = f.read() - size = len(contents) - - logging.debug(f"name: {name}") - logging.debug("position: ???") - logging.debug(f"file size: {size}") - binary_files.append((name, size, contents)) - - file_count += 1 - total_size += size - - binary_files.sort(key=lambda x: x[0]) - return binary_files, total_size, file_count - - def _create_file_list(self): - """Re-gather the necessary information on each file in the archive - while taking into account the modifications made by the user since - then. - """ - binary_files = [] - file_count = 0 - total_size = 0 - - for name in self.entries: - if name in self.modified_entries: - entry = self.modified_entries[name] - entry_bytes = entry.content - - if entry.action is FileAction.REMOVE: - logging.info(f"removing {name}") - continue - - logging.info(f"editing {name}") - else: - entry = self.entries[name] - entry_bytes = self._get_file(name) - - binary_files.append((entry.name, len(entry_bytes))) - file_count += 1 - total_size += len(entry_bytes) - - for entry in [x for x in self.modified_entries.values() if x.action is FileAction.ADD]: - logging.info(f"adding {entry.name}") - entry_bytes = entry.content - binary_files.append((entry.name, len(entry_bytes))) - file_count += 1 - total_size += len(entry_bytes) - - binary_files.sort(key=lambda x: x[0]) - return binary_files, total_size, file_count - - def _get_file(self, name: str) -> bytes: - """Get the contents of a specific file in the big based on file name""" - entry = self.entries[name] - with open(self.file_path, "rb") as f: - f.seek(entry.position) - return f.read(entry.size) - - def save(self, path: str = None): - """Save the archive to a file. - - Params - ------- - path : Optional[str] - The new path to save to. Something like 'path/to/file/test.big'. - Omit this if you just want to save in the same file. - """ - self._pack(path) - - @classmethod - def from_directory(cls, path: str, file_path: str) -> "LargeArchive": - """Generate a BIG archive from a directory. This is useful for - compiling an archive without adding each file manually. You simply - give the top level directory and every file will be added recursively. - - Params - ------- - path : str - Path to the top level folder of the files you wish to compile - file_path : str - Path to save the new archive - - Returns - -------- - Archive - Compiled archived - """ - raise NotImplementedError - - binary_files, total_size, file_count = cls._create_file_list_from_directory(path) - entries = cls._pack_file_list_from_directory(binary_files, total_size, file_count) - - return cls(file_path, entries=entries) - - @classmethod - def empty(cls, file_path): - """Generate an empty archive. - - Params - ------- - file_path : str - Path to save the archive to once it becomes used. - - """ - if os.path.exists(file_path): - raise ValueError(f"File {file_path} already exists.") - - with open(file_path, "wb") as f: - f.write(b"") - - return cls(file_path, entries={}) - - def repack(self): - """Update the archive to include all the modified entries. This clears - the list and updates the archive with the new data. - """ - self._pack(None) - - def archive_memory_size(self) -> int: - """Get the current in memory size of all the modifies entries that - have not yet been saved. You can use this to decide when you would - like to save in relation to the capacities of your machine. - """ - return sum( - [ - len(entry.content) - for entry in self.modified_entries.values() - if entry.action in (FileAction.ADD, FileAction.EDIT) - ] - ) diff --git a/pyBIG/memory_archive.py b/pyBIG/memory_archive.py new file mode 100644 index 0000000..f8aee05 --- /dev/null +++ b/pyBIG/memory_archive.py @@ -0,0 +1,132 @@ +import io +import logging +from typing import IO, Type, TypeVar + +from .base_archive import BaseArchive, FileList + +T = TypeVar("T", bound="InMemoryArchive") + + +class InMemoryArchive(BaseArchive): + """The core of the library, represents a BIG file and allows + the user to mainpulate it programatically. + + This implementation stores as much as possible in memory to avoid + creating unecessary clutter that needs to be cleaned up. All + disk action need to be explicit through .save + + Params + ------- + content : Optional[bytes] + Raw bytes of the original big file + + """ + + def __init__(self, content: bytes = b"", **kwargs): + self.archive = io.BytesIO(content) + self.entries = kwargs.get("entries") + self.modified_entries = {} + self.header = kwargs.get("header", "BIG4") + + if self.entries is None: + self.entries, self.header = self._unpack(self.archive) + + def __repr__(self): + return f"< Archive entries={len(self.entries)} dirty={bool(self.modified_entries)} >" + + def _pack(self): + """Rewrite the archive with the modifications stored + in self.modified_entries.""" + new_archive = io.BytesIO() + + file_data = self._create_file_list() + self.entries = self._pack_file_list(new_archive, *file_data, self.header) + + self._pack_files(new_archive, *file_data) + + self.archive = new_archive + self.archive.seek(0) + self.modified_entries = {} + + def _create_entry(self, name: str, content: bytes) -> tuple: + """In this in-memory archive, an entry is the name, size + of the file and the contents of the file + """ + + return name, len(content), content + + def _pack_files( + self, raw_data_file: IO, file_list: FileList, total_size: int, file_count: int + ): + """Combine all files into a single raw data bundle""" + + logging.info("packing files") + for file in file_list: + raw_data_file.write(file[2]) + logging.info("finished packing files") + + def _get_file(self, name: str) -> bytes: + """Get the contents of a specific file in the big based on file name""" + entry = self.entries[name] + self.archive.seek(entry.position) + return self.archive.read(entry.size) + + def save(self, path: str): + """Save the archive to a file. + + Params + ------- + path : str + The path to save to. Something like 'path/to/file/test.big' + """ + self._pack() + with open(path, "wb") as f: + f.write(self.archive.getvalue()) + + @classmethod + def from_directory(cls: Type[T], path: str, header: str = "BIG4") -> T: + """Generate a BIG archive from a directory. This is useful for + compiling an archive without adding each file manually. You simply + give the top level directory and every file will be added recursively. + + Params + ------- + path : str + Path to the top level folder of the files you wish to compile + header : str + The type of archive, either BIG4 or BIGF. Defaults to BIG4 + + Returns + -------- + Archive + Compiled archived + """ + return cls._pack_archive_from_directory(cls.empty(header), path) + + @classmethod + def empty(cls: Type[T], header: str = "BIG4") -> T: + """Generate an empty archive. + + Params + ------- + header : str + The type of the archive, can either be BIG4 or BIGF. Defaults to BIG4 + + Returns + -------- + Archive + Empty archive + """ + return cls(entries={}, header=header) + + def bytes(self): + """Returns the archive data as bytes + + Returns + -------- + bytes + The archive data + """ + self._pack() + + return self.archive.getvalue() diff --git a/pyBIG/refpack.py b/pyBIG/refpack.py new file mode 100644 index 0000000..d1aac82 --- /dev/null +++ b/pyBIG/refpack.py @@ -0,0 +1,271 @@ +import logging +import struct + + +def has_refpack_header(data: bytes) -> bool: + """Check if input data is refpack by checking the + header. Data may be headerless refpack data, only way + to really check is to attempt a decompress + + Params + ------- + data: bytes + Input data to check + + Returns + -------- + bool + True if data has the refpack header + """ + if len(data) < 5: + return False + + # Check 2-byte magic number + marker = struct.unpack(">H", data[:2])[0] + if marker != 0x10FB: + return False + + # Read 3-byte uncompressed size + size_bytes = data[2:5] + unpacked_size = size_bytes[0] | (size_bytes[1] << 8) | (size_bytes[2] << 16) + + # Basic sanity check + if unpacked_size == 0 or unpacked_size > 100_000_000: + return False + + return True + + +def matchlen(s: bytes, d: bytes, maxmatch: int) -> int: + current = 0 + while current < maxmatch and s[current] == d[current]: + current += 1 + return current + + +def hash_bytes(data: bytes) -> int: + return ((data[0] << 4) ^ (data[1] << 2) ^ (data[2])) & 0xFFFF + + +def compress(input_data: bytes) -> bytes: + """Compress bytes to refpack format + + Params + ------- + input_data: bytes + Data to compress + + + Returns + -------- + bytes + Compressed data + """ + length = len(input_data) + to = bytearray() + + # Add RefPack magic number (0x10FB) and uncompressed size (3 bytes LE) + to += struct.pack(">H", 0x10FB) # big-endian magic + to += bytes([(length >> 16) & 0xFF, (length >> 8) & 0xFF, length & 0xFF]) + + compressed = bytearray() + + run = 0 + cptr = 0 + rptr = 0 + + hashtbl = [-1] * 65536 + link = [-1] * 131072 + + while cptr < length: + boffset = 0 + blen = 2 + bcost = 2 + mlen = min(length - cptr, 1028) + if cptr + 2 >= length: + mlen = 0 + + if mlen >= 3: + h = hash_bytes(input_data[cptr : cptr + 3]) + hoffset = hashtbl[h] + minhoffset = max(cptr - 131071, 0) + + while hoffset >= minhoffset: + tptr = hoffset + if ( + cptr + blen < length + and tptr + blen < length + and input_data[cptr + blen] == input_data[tptr + blen] + ): + tlen = matchlen(input_data[cptr:], input_data[tptr:], mlen) + if tlen > blen: + toffset = (cptr - 1) - tptr + if toffset < 1024 and tlen <= 10: + tcost = 2 + elif toffset < 16384 and tlen <= 67: + tcost = 3 + else: + tcost = 4 + + if tlen - tcost + 4 > blen - bcost + 4: + blen = tlen + bcost = tcost + boffset = toffset + if blen >= 1028: + break + hoffset = link[hoffset & 131071] + + if bcost >= blen: + h = hash_bytes(input_data[cptr : cptr + 3]) if cptr + 2 < length else 0 + hoffset = cptr + link[hoffset & 131071] = hashtbl[h] + hashtbl[h] = hoffset + + run += 1 + cptr += 1 + else: + while run > 3: + tlen = min(112, run & ~3) + run -= tlen + compressed.append(0xE0 + (tlen >> 2) - 1) + compressed += input_data[rptr : rptr + tlen] + rptr += tlen + + if bcost == 2: + compressed.append(((boffset >> 8) << 5) + ((blen - 3) << 2) + run) + compressed.append(boffset & 0xFF) + elif bcost == 3: + compressed.append(0x80 + (blen - 4)) + compressed.append((run << 6) + (boffset >> 8)) + compressed.append(boffset & 0xFF) + else: + compressed.append(0xC0 + ((boffset >> 16) << 4) + (((blen - 5) >> 8) << 2) + run) + compressed.append((boffset >> 8) & 0xFF) + compressed.append(boffset & 0xFF) + compressed.append((blen - 5) & 0xFF) + + if run: + compressed += input_data[rptr : rptr + run] + rptr += run + run = 0 + + for i in range(blen): + if cptr + 2 < length: + h = hash_bytes(input_data[cptr : cptr + 3]) + hoffset = cptr + link[hoffset & 131071] = hashtbl[h] + hashtbl[h] = hoffset + cptr += 1 + + rptr = cptr + + while run > 3: + tlen = min(112, run & ~3) + run -= tlen + compressed.append(0xE0 + (tlen >> 2) - 1) + compressed += input_data[rptr : rptr + tlen] + rptr += tlen + + compressed.append(0xFC + run) + if run: + compressed += input_data[rptr : rptr + run] + + to += compressed + return bytes(to) + + +def decompress(input_data: bytes, ignore_mismatch: bool = False) -> bytes: + """Decompress refpack data. This expects the data to have a refpack header + but will still attempt to decompress if it cannot find + + Params + ------- + input_data: bytes + The data to decompress + ingore_mismatch: Optional[bool] + If the data has a refpack header, the function will + raise an error if the expected size is not the same + as the decompressed size. You can use this to suppres it. + + Returns + -------- + bytes + The decompressed bytes + """ + index = 0 + + expected_size = None + if len(input_data) >= 5: + magic = struct.unpack(">H", input_data[:2])[0] + if magic == 0x10FB: + expected_size = (input_data[2] << 16) | (input_data[3] << 8) | input_data[4] + input_data = input_data[5:] + + output = bytearray() + + while True: + first = input_data[index] + index += 1 + + if not (first & 0x80): # short ref + second = input_data[index] + index += 1 + run = first & 3 + output += input_data[index : index + run] + index += run + ref_offset = ((first & 0x60) << 3) + second + ref = len(output) - 1 - ref_offset + length_to_copy = ((first & 0x1C) >> 2) + 3 + for _ in range(length_to_copy): + output.append(output[ref]) + ref += 1 + continue + + if not (first & 0x40): # long ref + second = input_data[index] + third = input_data[index + 1] + index += 2 + run = second >> 6 + output += input_data[index : index + run] + index += run + ref_offset = ((second & 0x3F) << 8) + third + ref = len(output) - 1 - ref_offset + length_to_copy = (first & 0x3F) + 4 + for _ in range(length_to_copy): + output.append(output[ref]) + ref += 1 + continue + + if not (first & 0x20): # very long ref + second = input_data[index] + third = input_data[index + 1] + fourth = input_data[index + 2] + index += 3 + run = first & 3 + output += input_data[index : index + run] + index += run + ref_offset = ((first & 0x10) >> 4 << 16) + (second << 8) + third + ref = len(output) - 1 - ref_offset + length_to_copy = (((first & 0x0C) >> 2) << 8) + fourth + 5 + for _ in range(length_to_copy): + output.append(output[ref]) + ref += 1 + continue + + # literal or EOF + run = ((first & 0x1F) << 2) + 4 + if run <= 112: + output += input_data[index : index + run] + index += run + continue + run = first & 3 + output += input_data[index : index + run] + break + + if expected_size is not None and expected_size != len(output): + if ignore_mismatch is True: + logging.info("Decompress size mismatch") + else: + raise ValueError("Decompress size mismatch") + + return bytes(output) diff --git a/setup.py b/setup.py index 347a845..2113a3c 100644 --- a/setup.py +++ b/setup.py @@ -20,6 +20,7 @@ description="A library for manipulating BIG format archives", long_description_content_type="text/markdown", long_description=readme, + python_requires='>=3.8', classifiers=[ "License :: OSI Approved :: MIT License", "Intended Audience :: Developers", @@ -28,6 +29,8 @@ "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "Topic :: Software Development :: Libraries", "Topic :: Software Development :: Libraries :: Python Modules", "Topic :: Utilities", diff --git a/tests/README.md b/tests/README.md deleted file mode 100644 index eb63908..0000000 --- a/tests/README.md +++ /dev/null @@ -1,6 +0,0 @@ -# Tests - -Tests must be run from root directory -* `python -m unittest tests.functional_tests` -* `python -m unittest tests.memory_tests` -* `python -m unittest tests.profiler` diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/functional_tests.py b/tests/functional_tests.py index 5d5c8c8..5d7ba13 100644 --- a/tests/functional_tests.py +++ b/tests/functional_tests.py @@ -1,8 +1,12 @@ +import logging import os from typing import Union import unittest -from pyBIG import Archive, LargeArchive +from pyBIG import InMemoryArchive, InDiskArchive +from pyBIG.refpack import compress, decompress, has_refpack_header + +logging.basicConfig(level=logging.INFO) TEST_FILE = "read_me_for_test.txt" @@ -13,7 +17,19 @@ class BaseTestCases: class BaseTest(unittest.TestCase): - archive: Union[Archive, LargeArchive] + archive: Union[InMemoryArchive, InDiskArchive] + + def tearDown(self): + for file in [ + f"tests/test_data/output/{TEST_FILE}", + "tests/test_data/output/test.big", + TEST_ARCHIVE, + "tests/test_data/test_big_type.big", + ]: + try: + os.remove(file) + except OSError: + pass def test_decode(self): contents = self.archive.read_file(TEST_FILE) @@ -54,35 +70,59 @@ def test_extract_and_load(self): self.archive.extract("tests/test_data/output") self.assertTrue(os.path.exists(f"tests/test_data/output/{TEST_FILE}")) - new_archive = Archive.from_directory("tests/test_data/output") + new_archive = InMemoryArchive.from_directory("tests/test_data/output") self.assertIn(TEST_FILE, new_archive.entries) new_archive.save("tests/test_data/output/test.big") self.assertTrue(os.path.exists("tests/test_data/output/test.big")) - os.remove(f"tests/test_data/output/{TEST_FILE}") - os.remove("tests/test_data/output/test.big") - def test_utils(self): self.archive.file_list() + def test_archive_memory_size(self): + archive = InDiskArchive.empty(file_path=TEST_ARCHIVE) + file_bytes = TEST_CONTENT.encode(TEST_ENCODING) + archive.add_file(TEST_FILE, file_bytes) + + size = archive.archive_memory_size() + + self.assertEqual(size, len(file_bytes)) + + def test_archive_bytes(self): + data = self.archive.bytes() + self.assertIsInstance(data, bytes) + class TestArchive(BaseTestCases.BaseTest): def setUp(self): with open("tests/test_data/test_big.big", "rb") as f: - self.archive = Archive(f.read()) + self.archive = InMemoryArchive(f.read()) def test_empty_archive(self): - archive = Archive.empty() + archive = InMemoryArchive.empty() archive.add_file(TEST_FILE, TEST_CONTENT.encode(TEST_ENCODING)) archive.repack() self.assertIn(TEST_FILE, archive.entries) + def test_archive_type(self): + path = "tests/test_data/test_big_type.big" + + for header in ["BIG4", "BIGF"]: + archive = InMemoryArchive.empty(header) + archive.add_file(TEST_FILE, TEST_CONTENT.encode(TEST_ENCODING)) + archive.save(path) + + with open(path, "rb") as f: + archive = InMemoryArchive(f.read()) + + self.assertEqual(archive.header, header) + os.remove("tests/test_data/test_big_type.big") + class TestLargeArchive(BaseTestCases.BaseTest): def setUp(self): - self.archive = LargeArchive("tests/test_data/test_big.big") + self.archive = InDiskArchive("tests/test_data/test_big.big") if self.archive.file_exists(TEST_FILE): self.archive.remove_file(TEST_FILE) @@ -94,22 +134,72 @@ def setUp(self): self.archive.repack() def test_empty_archive(self): - archive = LargeArchive.empty(TEST_ARCHIVE) + archive = InDiskArchive.empty(file_path=TEST_ARCHIVE) archive.add_file(TEST_FILE, TEST_CONTENT.encode(TEST_ENCODING)) archive.repack() self.assertIn(TEST_FILE, archive.entries) - os.remove(TEST_ARCHIVE) - - def test_archive_memory_size(self): - archive = LargeArchive.empty(TEST_ARCHIVE) - file_bytes = TEST_CONTENT.encode(TEST_ENCODING) - archive.add_file(TEST_FILE, file_bytes) - - size = archive.archive_memory_size() - self.assertEqual(size, len(file_bytes)) - os.remove(TEST_ARCHIVE) + def test_archive_type(self): + path = "tests/test_data/test_big_type.big" + + for header in ["BIG4", "BIGF"]: + archive = InDiskArchive.empty(header, file_path=path) + archive.add_file(TEST_FILE, TEST_CONTENT.encode(TEST_ENCODING)) + archive.save() + archive = InDiskArchive(path) + self.assertEqual(archive.header, header) + os.remove("tests/test_data/test_big_type.big") + + +class TestRefPack(unittest.TestCase): + def test_refpack_check_valid_data(self): + data = b"Sample data for testing." + compressed = compress(data) + self.assertTrue(has_refpack_header(compressed)) + + def test_refpack_check_invalid_magic(self): + invalid_data = b"\x00\x00\x00\x00" + self.assertFalse(has_refpack_header(invalid_data)) + + def test_refpack_check_invalid_size(self): + # Construct data with valid magic but invalid size (0) + data = b"\x10\xfb\x00\x00" + self.assertFalse(has_refpack_header(data)) + + def test_compress_decompress_roundtrip(self): + data = b"Example data to compress and decompress." + compressed = compress(data) + decompressed = decompress(compressed) + self.assertEqual(decompressed, data) + + def test_decompress_size_mismatch_raises(self): + data = b"Test data with mismatch" + compressed = compress(data) + corrupted = bytearray(compressed) + corrupted[2:5] = b"\x00\x00\x00" # corrupt expected size + with self.assertRaises(ValueError): + decompress(bytes(corrupted), ignore_mismatch=False) + + def test_decompress_size_mismatch_ignore(self): + data = b"Test data with mismatch" + compressed = compress(data) + corrupted = bytearray(compressed) + corrupted[2:5] = b"\x00\x00\x00" + decompressed = decompress(bytes(corrupted), ignore_mismatch=True) + self.assertTrue(decompressed) # decompress returns bytes even if size mismatch ignored + + def test_empty_data(self): + data = b"" + compressed = compress(data) + decompressed = decompress(compressed) + self.assertEqual(decompressed, data) + + def test_large_data(self): + data = b"A" * 10_000 + compressed = compress(data) + decompressed = decompress(compressed) + self.assertEqual(decompressed, data) if __name__ == "__main__": diff --git a/tests/memory_tests.py b/tests/memory_tests.py index 8f6aa96..4243c33 100644 --- a/tests/memory_tests.py +++ b/tests/memory_tests.py @@ -1,18 +1,18 @@ import os import unittest -from pyBIG import Archive, LargeArchive +from pyBIG import InMemoryArchive, InDiskArchive from pyBIG.utils import getsize class MemoryTests(unittest.TestCase): def test_1(self): - archive = Archive.empty() + archive = InMemoryArchive.empty() archive.add_file("test_file.txt", b"Clement1" * 1250000) in_memory_size = getsize(archive) - archive = LargeArchive.empty("big_big.big") + archive = InDiskArchive.empty(file_path="big_big.big") archive.add_file("test_file.txt", b"Clement1" * 1250000) archive.save() @@ -23,13 +23,13 @@ def test_1(self): assert in_memory_size > on_disk_size def test_2(self): - archive = LargeArchive.empty("big_big.big") + archive = InDiskArchive.empty(file_path="big_big.big") archive.add_file("test_file.txt", b"Clement1" * 1250000) archive.save() post_save_size = getsize(archive) - archive = LargeArchive("big_big.big") + archive = InDiskArchive("big_big.big") loaded_size = getsize(archive) diff --git a/tests/profiler.py b/tests/profiler.py index 1319a8a..482b0d6 100644 --- a/tests/profiler.py +++ b/tests/profiler.py @@ -9,7 +9,7 @@ def speed(): with open("test_data/__edain_data.big", "rb") as f: - archive = pyBIG.Archive(f.read()) + archive = pyBIG.InMemoryArchive(f.read()) archive.add_file("weapons.ini", b"") pr = cProfile.Profile() @@ -30,9 +30,9 @@ def speed(): def size(): with open("test_data/__edain_data.big", "rb") as f: - archive = pyBIG.Archive(f.read()) + archive = pyBIG.InMemoryArchive(f.read()) - large_archive = pyBIG.LargeArchive("test_data/__edain_data.big") + large_archive = pyBIG.InDiskArchive("test_data/__edain_data.big") logging.info(f"Archive: {objsize.get_deep_size(archive)}") logging.info(f"Large Archive: {objsize.get_deep_size(large_archive)}")