itar 0.1.0__py3-none-any.whl

itar/__init__.py

@@ -0,0 +1,5 @@
+from .sharded_indexed_tar import ShardedIndexedTar
+
+__all__ = [
+    "ShardedIndexedTar",
+]

itar/sharded_indexed_tar.py

@@ -0,0 +1,256 @@
+import os
+from collections.abc import Mapping
+from io import BufferedReader
+from pathlib import Path
+from tarfile import TarInfo
+from typing import IO, Callable
+
+from .utils import (
+    TarFileSectionIO,
+    TarIndexError,
+    TarMember,
+    ThreadSafeFileIO,
+    build_tar_index,
+    check_tar_index,
+    tar_file_info,
+)
+
+ShardedTarIndex = dict[str, (int, TarMember)]  # fname -> (shard_idx, TarMember)
+
+
+class ShardedIndexedTar(Mapping):
+    def __init__(
+        self,
+        shards: list[str | os.PathLike | IO[bytes]],
+        index: ShardedTarIndex | None = None,
+        open_fn: Callable[[str | os.PathLike], IO[bytes]] = None,
+        buffered_file_reader: bool = True,
+        progress_bar: bool = False,
+    ):
+        """
+        For best performance, we recommend using unbuffered shard IO in combination with buffered individual file section readers.
+        In our (limited) benchmarks, this was the fastest configuration for reading many small files.
+        It it possbile that in other scenarios, buffered shard IO and unbuffered file section readers are faster.
+        """
+        self._needs_open = [isinstance(s, str | os.PathLike) for s in shards]
+        self._file_reader = BufferedReader if buffered_file_reader else lambda x: x
+        open_fn = (
+            open_fn or ThreadSafeFileIO
+        )  # In our benchmarks, `ThreadSafeFileIO` is even faster than `partial(open, mode="rb", buffering=0)`. Likely due to `pread` being fewer syscalls than `seek` + `read`.
+        self._shard_file_objs: IO[bytes] = [
+            open_fn(tar) if needs_open else tar
+            for tar, needs_open in zip(shards, self._needs_open, strict=True)
+        ]
+        if progress_bar:
+            from tqdm import tqdm
+        else:
+            tqdm = lambda x, **kwargs: x  # noqa: E731
+        self._index = (
+            index
+            if index is not None
+            else {
+                name: (i, member)
+                for i, file_obj in enumerate(
+                    tqdm(self._shard_file_objs, desc="Building index", unit="shard")
+                )
+                for name, member in build_tar_index(file_obj).items()
+            }
+        )
+
+    @classmethod
+    def open(
+        cls,
+        path: str | os.PathLike,
+        shards: list[str | os.PathLike] | None = None,
+        open_fn: Callable[[str | os.PathLike], IO[bytes]] = None,
+        buffered_file_reader: bool = True,
+    ):
+        import msgpack
+
+        path = Path(path)
+        with open(path, "rb") as f:
+            num_shards, index = msgpack.load(f)
+        return cls(
+            shards
+            if shards is not None
+            else [cls.shard_path(path, num_shards, i) for i in range(num_shards)],
+            index,
+            open_fn=open_fn,
+            buffered_file_reader=buffered_file_reader,
+        )
+
+    def save(self, path: str | os.PathLike):
+        import msgpack
+
+        path = Path(path)
+        with open(path, "wb") as f:
+            msgpack.dump((len(self._shard_file_objs), self.index), f)
+
+    @staticmethod
+    def shard_path(path: str | os.PathLike, num_shards: int, shard_idx: int) -> Path:
+        path = Path(path)
+        return path.parent / f"{path.stem}-{shard_idx:0{len(str(num_shards - 1))}d}.tar"
+
+    def file(self, name: str) -> IO[bytes]:
+        i, member = self._index[name]
+        _, offset_data, size = member
+        if isinstance(size, str):
+            return self.file(size)  # symlink or hard link
+        return self._file_reader(
+            TarFileSectionIO(self._shard_file_objs[i], offset_data, size)
+        )
+
+    def info(self, name: str) -> TarInfo:
+        i, member = self._index[name]
+        offset, _, _ = member
+        return tar_file_info(offset, self._shard_file_objs[i])
+
+    def check_tar_index(self, names: list[str] | None = None):
+        for name in names if names is not None else self:
+            i, member = self._index[name]
+            check_tar_index(name, member, self._shard_file_objs[i])
+
+    @property
+    def index(self) -> ShardedTarIndex:
+        return self._index
+
+    def close(self):
+        for needed_open, file_obj in zip(
+            self._needs_open, self._shard_file_objs, strict=True
+        ):
+            if needed_open:
+                # only close what we opened
+                file_obj.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
+    def __getitem__(self, name: str):
+        return self.file(name)
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._index
+
+    def __iter__(self):
+        return iter(self._index)
+
+    def __len__(self):
+        return len(self._index)
+
+    def keys(self):
+        return self._index.keys()
+
+    def values(self):
+        for name in self._index:
+            yield self[name]
+
+    def items(self):
+        for name in self._index:
+            yield (name, self[name])
+
+
+def cli():
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Create or check an itar index.")
+    subparsers = parser.add_subparsers(dest="command")
+
+    create_parser = subparsers.add_parser("create", help="Create an itar index")
+    create_parser.add_argument(
+        "itar", type=Path, help="Path to the new itar index that will be created"
+    )
+
+    check_parser = subparsers.add_parser("check", help="Check an existing itar index")
+    check_parser.add_argument("itar", type=Path, help="Path to the itar index")
+
+    ls_parser = subparsers.add_parser("ls", help="List files in an itar index")
+    ls_parser.add_argument("itar", type=Path, help="Path to the itar index")
+    ls_parser.add_argument(
+        "-l", "--long", action="store_true", help="Show long listing format"
+    )
+    ls_parser.add_argument(
+        "-H", "--human-readable", action="store_true", help="Use human-readable sizes"
+    )
+
+    args = parser.parse_args()
+
+    if args.command == "create":
+        _create(args)
+    elif args.command == "check":
+        _check(args)
+    elif args.command == "ls":
+        _ls(args)
+
+
+def _create(args):
+    import re
+
+    itar_path = Path(args.itar)
+    stem = itar_path.stem
+    pattern = re.compile(rf"^{re.escape(stem)}-\d+\.tar$")
+    shards = sorted(
+        [
+            s
+            for s in itar_path.parent.glob(f"{itar_path.stem}-*.tar")
+            if s.is_file() and pattern.match(s.name)
+        ]
+    )
+    num_shards = len(shards)
+    if num_shards < 1:
+        print(
+            f"No shards found for {itar_path}.\n"
+            f"Please create shard files first. Expected pattern: "
+            f"'{itar_path.stem}-NN.tar' (where NN is the zero-padded shard index, starting from 0)."
+        )
+        exit(1)
+
+    # ensure shards are named correctly
+    expected = [
+        ShardedIndexedTar.shard_path(itar_path, num_shards, i)
+        for i in range(num_shards)
+    ]
+    assert shards == expected, (
+        f"Shards do not match expected names: {shards} != {expected}"
+    )
+
+    with ShardedIndexedTar(shards, progress_bar=True) as itar:
+        itar.save(itar_path)
+    print(f"Created itar index at {itar_path} with {num_shards} shards.")
+
+
+def _check(args):
+    from tqdm import tqdm
+
+    did_error = False
+
+    with ShardedIndexedTar.open(args.itar) as itar:
+        for member in tqdm(itar, desc="Checking files", unit="file"):
+            try:
+                itar.check_tar_index([member])
+            except TarIndexError as e:
+                print(e)
+                did_error = True
+
+    if did_error:
+        exit(1)
+
+
+def _ls(args):
+    with ShardedIndexedTar.open(args.itar) as itar:
+        index = itar.index
+    if args.long:
+        for member, (shard_idx, (offset, offset_data, size)) in index.items():
+            if args.human_readable:
+                from humanize import naturalsize
+
+                size = naturalsize(size, gnu=True)
+            print(
+                f"{member:<40} {shard_idx:>5} {offset:>12} {offset_data:>12} {size:>10}"
+            )
+        print(f"{'NAME':<40} {'SHARD':>5} {'OFFSET':>12} {'OFF_DATA':>12} {'SIZE':>10}")
+    else:
+        for member in index:
+            print(member)

itar/utils.py

@@ -0,0 +1,230 @@
+import io
+import os
+import tarfile
+import threading
+from contextlib import nullcontext
+from tarfile import TarFile, TarInfo
+from types import SimpleNamespace
+from typing import IO, BinaryIO
+
+TarMember = tuple[int, int, int | str]  # (offset, offset_data, size | linkname)
+
+
+def tar_file_info(offset: int, file_obj: IO[bytes]) -> TarInfo:
+    file_obj.seek(offset)
+    return TarInfo.fromtarfile(
+        # want to avoid creating a new TarFile instance (potentially slow)
+        SimpleNamespace(
+            fileobj=file_obj,
+            # would be the defaults after TarFile.__init__
+            encoding=TarFile.encoding,
+            errors="surrogateescape",
+            pax_headers={},
+        )
+    )
+
+
+def tarinfo2member(tarinfo: TarInfo) -> TarMember:
+    if tarinfo.issym():
+        size = os.path.normpath(
+            "/".join(filter(None, (os.path.dirname(tarinfo.name), tarinfo.linkname)))
+        )
+    elif tarinfo.islnk():
+        size = os.path.normpath(tarinfo.linkname)
+    else:
+        size = tarinfo.size
+
+    if tarinfo.sparse is not None:
+        raise NotImplementedError("Sparse files are not supported")
+
+    return (tarinfo.offset, tarinfo.offset_data, size)
+
+
+def build_tar_index(
+    tar: str | os.PathLike | IO[bytes] | TarFile,
+) -> dict[str, TarMember]:
+    if isinstance(tar, str | os.PathLike):
+        tar = tarfile.open(tar, "r:")
+    elif isinstance(tar, TarFile):
+        tar = nullcontext(tar)
+    else:
+        tar.seek(0)
+        tar = tarfile.open(fileobj=tar, mode="r:")
+
+    with tar as f:
+        members = {member.name: member for member in f.getmembers()}
+        return {
+            member.name: tarinfo2member(member)
+            for member in members.values()
+            if (
+                # index only includes files and links. no directories, devices, etc.
+                member.isreg()
+                or (
+                    member.type not in tarfile.SUPPORTED_TYPES
+                )  # Members with unknown types are treated as regular files.
+                or member.issym()
+                or member.islnk()
+            )
+        }
+
+
+class TarIndexError(Exception):
+    pass
+
+
+def check_tar_index(
+    name: str,
+    tar_offset: TarMember,
+    file_obj: IO[bytes],
+):
+    offset, offset_data, size = tar_offset
+    info = tar_file_info(offset, file_obj)
+    if (
+        info.offset != offset
+        or info.name != name
+        or info.offset_data != offset_data
+        or (info.size != size and not info.islnk() and not info.issym())
+    ):
+        # TODO: if link, check that linkname is correct
+        raise TarIndexError(
+            f"Index mismatch: "
+            f"expected ({name}, {offset}, {offset_data}, {size}), "
+            f"got ({info.name}, {info.offset}, {info.offset_data}, {info.size})"
+        )
+
+
+class ThreadSafeFileIO(io.RawIOBase):
+    """
+    A thread-safe, file-like object that wraps a file descriptor
+    and uses os.pread() for concurrent reads.
+
+    Each thread has its own seek position.
+    """
+
+    def __init__(self, path: str | os.PathLike):
+        self._path = str(path)
+        self._fd = os.open(path, os.O_RDONLY)
+        self._local = threading.local()
+        self._closed = False
+
+    def _get_pos(self) -> int:
+        return getattr(self._local, "pos", 0)
+
+    def _set_pos(self, val: int) -> None:
+        self._local.pos = val
+
+    def seek(self, offset: int, whence: int = os.SEEK_SET) -> int:
+        pos = self._get_pos()
+        if whence == os.SEEK_SET:
+            self._set_pos(offset)
+        elif whence == os.SEEK_CUR:
+            self._set_pos(pos + offset)
+        elif whence == os.SEEK_END:
+            end = os.lseek(self._fd, 0, os.SEEK_END)
+            self._set_pos(end + offset)
+        else:
+            raise ValueError(f"Invalid whence: {whence}")
+        return self._get_pos()
+
+    def tell(self) -> int:
+        return self._get_pos()
+
+    def read(self, size: int = -1) -> bytes:
+        if size < 0:
+            raise NotImplementedError("read(size=-1) not supported")
+        pos = self._get_pos()
+        data = os.pread(self._fd, size, pos)
+        self._set_pos(pos + len(data))
+        return data
+
+    def readable(self) -> bool:
+        return True
+
+    def seekable(self) -> bool:
+        return True
+
+    def writable(self) -> bool:
+        return False
+
+    def close(self) -> None:
+        if not self._closed:
+            os.close(self._fd)
+            self._closed = True
+
+    def fileno(self) -> int:
+        return self._fd
+
+    @property
+    def name(self) -> str:  # optional
+        return self._path
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, tb):
+        self.close()
+
+
+class TarFileSectionIO(io.RawIOBase):
+    def __init__(self, fileobj: BinaryIO, offset: int, size: int):
+        self._fileobj = fileobj
+        self._start = offset
+        self._end = offset + size
+        self._pos = 0  # relative position within the slice
+
+    def read(self, size: int = -1) -> bytes:
+        if self._pos >= self._end - self._start:
+            return b""
+
+        abs_pos = self._start + self._pos
+        max_len = self._end - (self._start + self._pos)
+
+        if size < 0 or size > max_len:
+            size = max_len
+
+        self._fileobj.seek(abs_pos)
+        data = self._fileobj.read(size)
+        self._pos += len(data)
+        return data
+
+    def readinto(self, b: bytearray) -> int:
+        data = self.read(len(b))
+        n = len(data)
+        b[:n] = data
+        return n
+
+    def readall(self) -> bytes:
+        return self.read(-1)
+
+    def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
+        if whence == io.SEEK_SET:
+            new_pos = offset
+        elif whence == io.SEEK_CUR:
+            new_pos = self._pos + offset
+        elif whence == io.SEEK_END:
+            new_pos = (self._end - self._start) + offset
+        else:
+            raise ValueError(f"Invalid whence: {whence}")
+
+        if new_pos < 0:
+            raise ValueError("Negative seek position")
+        self._pos = min(new_pos, self._end - self._start)
+        return self._pos
+
+    def tell(self) -> int:
+        return self._pos
+
+    def readable(self) -> bool:
+        return True
+
+    def seekable(self) -> bool:
+        return True
+
+    def writable(self) -> bool:
+        return False
+
+    def close(self) -> None:
+        pass  # Don't close the underlying file
+
+    def __len__(self) -> int:
+        return self._end - self._start

itar-0.1.0.dist-info/METADATA

@@ -0,0 +1,30 @@
+Metadata-Version: 2.3
+Name: itar
+Version: 0.1.0
+Summary: Add your description here
+Author: Karim Abou Zeid
+Author-email: Karim Abou Zeid <contact@ka.codes>
+Requires-Dist: humanize>=4.12.3
+Requires-Dist: msgpack>=1.1.1
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# itar
+
+## Minimal example
+
+```bash
+echo Hello world! > bar.txt
+tar cf foo-0.tar bar.txt  # create regular tar file(s) with zero-padded shard number
+itar create foo.itar  # create the index
+itar ls foo.itar  # (optional) view the contents of the index
+```
+
+Now you can open and efficiently access the archived files in python in constant time
+```python
+from itar import ShardedIndexedTar
+
+with ShardedIndexedTar.open("foo.itar") as itar:
+    f = itar.file("bar.txt")  # file-like object
+    assert f.read().decode("utf-8") == "Hello world!\n"
+```

itar-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+itar/__init__.py,sha256=1G91ozs-YBv-lnHeStSoDk6HTye_HLkIrLn_fQjSVt4,91
+itar/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+itar/sharded_indexed_tar.py,sha256=G3k4kD4_5vakAzsBC7RfZ0YGPk2FuiDtmsE08fxAZG8,8141
+itar/utils.py,sha256=iduueK8wfnEBwd1Tcyt6z544gpJYA62YXPI9jdc1Xbc,6453
+itar-0.1.0.dist-info/WHEEL,sha256=I8-bO5cg2sb8TH6ZM6EgCP87Y1cV_f9UGgWnfAhVOZI,78
+itar-0.1.0.dist-info/entry_points.txt,sha256=uRZGXYbybs8I9NbVn6dODCzi30Z5BEEZ3b6bR93ot1E,55
+itar-0.1.0.dist-info/METADATA,sha256=x-IiG0GfVw5hoQ8JI2rGqslRIYNSuD-ygolTxAJAHL8,830
+itar-0.1.0.dist-info/RECORD,,

itar-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.9.0
+Root-Is-Purelib: true
+Tag: py3-none-any

itar-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+itar = itar.sharded_indexed_tar:cli
+