diff --git a/src/pymsi/package.py b/src/pymsi/package.py index 43ecaee..8249542 100644 --- a/src/pymsi/package.py +++ b/src/pymsi/package.py @@ -1,5 +1,6 @@ import copy import io +import mmap from pathlib import Path from typing import Iterator, Optional, Union @@ -18,13 +19,14 @@ class Package: - def __init__(self, path_or_bytesio: Union[Path, io.BytesIO]): - if isinstance(path_or_bytesio, io.BytesIO): - self.path = None - self.file = path_or_bytesio - else: + # TODO: consider typing.BinaryIO + def __init__(self, path_or_bytesio: Union[Path, io.BytesIO, mmap.mmap]): + if isinstance(path_or_bytesio, Path): self.path = path_or_bytesio.resolve(True) self.file = self.path.open("rb") + else: + self.path = None + self.file = path_or_bytesio self.tables = {} self.ole = None self.summary = None diff --git a/tests/io.py b/tests/io.py new file mode 100644 index 0000000..6636c2f --- /dev/null +++ b/tests/io.py @@ -0,0 +1,49 @@ +import mmap +from pathlib import Path + +import pytest + +import pymsi + + +def read_package(path_or_bytesio): + with pymsi.Package(path_or_bytesio) as package: + msi = pymsi.Msi(package) + return msi + + +# Function to read a package using Path +def read_package_path(file_path): + path = Path(file_path) + return read_package(path) + + +# Function to read a package using with open +def read_package_with_open(file_path): + with open(file_path, "rb") as f: + return read_package(f) + + +# Function to read a package using mmap +def read_package_mmap(file_path): + with open(file_path, "r+b") as f: + mm = mmap.mmap(f.fileno(), 0) + return read_package(mm) + + +# Test cases +@pytest.mark.parametrize( + "read_package_func", [read_package_path, read_package_with_open, read_package_mmap] +) +def test_read_package(read_package_func): + test_file = "powertoys.msi" + + msi = read_package_func(test_file) + + size = (msi.package.ole.nb_sect * msi.package.ole.sector_size) + 512 + + print(size) + + +if __name__ == "__main__": + pytest.main()