(shared.PakFile) fixed __init__ & .as_bytes() + added tests

bsp_tool/branches/shared.py

@@ -118,19 +118,28 @@ def from_bytes(cls, raw_lump: bytes):
 class PakFile(zipfile.ZipFile):
     _buffer: io.BytesIO
 
-    def __init__(self, file: Any = None):
+    def __init__(self, file_: Any = None, mode: str = "a", **kwargs):
         """always a read-only copy of the lump"""
-        if file is None:
-            self._buffer = io.BytesIO(b"")
-        elif isinstance(file, io.BytesIO):
-            self._buffer = file
-        elif isinstance(file, str):
-            self._buffer = io.BytesIO(open(file, "rb").read())
+        if file_ is None:
+            empty_zip = [b"PK\x05\x06", b"\x00" * 16, b"\x20\x00XZP1\x20\x30", b"\x00" * 26]
+            self._buffer = io.BytesIO(b"".join(empty_zip))
+        elif isinstance(file_, io.BytesIO):  # BspClass will take this route via .from_bytes()
+            self._buffer = file_
+        elif isinstance(file_, str):
+            self._buffer = io.BytesIO(open(file_, "rb").read())
         else:
-            raise TypeError(f"Cannot create {self.__class__.__name__} from type '{type(file)}'")
-        super().__init__(self._buffer, "r")
+            raise TypeError(f"Cannot create {self.__class__.__name__} from type '{type(file_)}'")
+        super().__init__(self._buffer, mode=mode, **kwargs)
 
     def as_bytes(self) -> bytes:
+        # write ending records if edits were made (adapted from ZipFile.close)
+        if self.mode in "wxa" and self._didModify and self.fp is not None:
+            with self._lock:
+                if self._seekable:
+                    self.fp.seek(self.start_dir)
+                self._write_end_record()
+        self._didModify = False  # don't double up when .close() is called
+        # NOTE: .close() can get funky but it's OK because ._buffer isn't a real file
         return self._buffer.getvalue()
 
     @classmethod

tests/branches/test_shared.py

@@ -0,0 +1,81 @@
+import os
+
+from bsp_tool.branches import shared
+
+import pytest
+
+
+class TestEntities:
+    # TODO: test all parser edge cases
+    ...
+
+
+class TestPakFile:
+    zips = {"empty.zip": b"".join([b"PK\x05\x06", b"\x00" * 16,
+                                   b"\x20\x00XZP1 0", b"\x00" * 26]),
+            "deflate.zip": b"".join([b"PK\x03\x04\x14", b"\x00" * 5,
+                                     b"\x92\x6E\xEF\x56\x99\xA0\xDC\x42",
+                                     b"\x07\x00\x00\x00" * 2, b"\x08\x00\x00\x00",
+                                     b"test.txthello~\n",
+                                     b"PK\x01\x02\x14\x03\x14", b"\x00" * 5,
+                                     b"\x92\x6E\xEF\x56\x99\xA0\xDC\x42",
+                                     b"\x07\x00\x00\x00" * 2, b"\x08\x00\x00\x00",
+                                     b"\x00" * 8, b"\x80\x01\x00\x00\x00\x00test.txt"
+                                     b"PK\x05\x06\x00\x00\x00\x00\x01\x00\x01\x00",
+                                     b"\x36\x00\x00\x00\x2D\x00\x00\x00",
+                                     b"\x20\x00XZP1 0", b"\x00" * 26])}
+    expected = {"empty.zip": dict(),
+                "deflate.zip": {"test.txt": b"hello~\n"}}
+
+    def test_new(self):
+        """create & populate a PakFile from nothing"""
+        pk = shared.PakFile()
+        pk.writestr("test.txt", "hello~\n")
+        assert pk.namelist() == ["test.txt"]
+        assert pk.read("test.txt") == b"hello~\n"
+
+    def setup_method(self, method):
+        """create test zipfiles"""
+        if method.__name__ in ("test_from_file", "test_bytes"):
+            for filename, data in self.zips.items():
+                with open(filename, "wb") as zip_file:
+                    zip_file.write(data)
+
+    def teardown_method(self, method):
+        """delete test zipfiles"""
+        if method.__name__ in ("test_from_file", "test_bytes"):
+            for filename in self.zips:
+                os.remove(filename)
+
+    @pytest.mark.parametrize("test_zip", zips)
+    def test_from_file(self, test_zip: str):
+        """open a .zip file"""
+        pk = shared.PakFile(test_zip)
+        assert set(pk.namelist()) == set(self.expected[test_zip])
+        for filename in pk.namelist():
+            assert pk.read(filename) == self.expected[test_zip][filename]
+
+    @pytest.mark.parametrize("test_zip", zips)
+    def test_bytes(self, test_zip: str):
+        raw_zip = self.zips[test_zip]
+        pk = shared.PakFile.from_bytes(raw_zip)
+        assert pk.as_bytes() == raw_zip
+
+    def test_save_changes(self):
+        pk = shared.PakFile()
+        pk.writestr("test.txt", "hello~\n")
+        raw_zip = pk.as_bytes()
+        # valid zip
+        pk2 = shared.PakFile.from_bytes(raw_zip)
+        assert pk.namelist() == pk2.namelist()
+        for filename in pk.namelist():
+            assert pk.read(filename) == pk2.read(filename)
+        # continue editing
+        pk.writestr("test2.txt", "~world\n")
+        # pk.close()
+        raw_zip = pk.as_bytes()
+        # valid zip
+        pk2 = shared.PakFile.from_bytes(raw_zip)
+        assert pk.namelist() == pk2.namelist()
+        for filename in pk.namelist():
+            assert pk.read(filename) == pk2.read(filename)