brkraw 0.3.11__py3-none-any.whl → 0.5.0__py3-none-any.whl

brkraw/core/fs.py

@@ -0,0 +1,495 @@
+"""
+Unified filesystem view for Paravision-like datasets stored as directories or zip
+archives.
+
+This module provides `DatasetFS`, a lightweight abstraction that presents the same
+API whether the dataset lives on disk or inside a zip file. It handles:
+- Anchor detection so zip members can be referenced with stable, anchor-stripped
+  relative paths.
+- Directory and zip traversal via a zipfile-like `walk` that mirrors `os.walk`.
+- Opening files by archive-relative path, yielding file-like objects or concrete
+  temp files for consumers that require real paths.
+- Repacking subtrees back into zip files with optional root folder control.
+
+`DatasetFS` is intentionally small and side-effect free so it can be reused
+outside Paravision-specific contexts.
+"""
+from __future__ import annotations
+from dataclasses import dataclass, field
+import io
+from pathlib import Path
+from typing import IO, Iterable, Literal, Optional, Tuple, List, Union, TYPE_CHECKING
+
+import os
+import zipfile
+import shutil
+from tempfile import TemporaryDirectory
+from . import zip as zipcore
+
+if TYPE_CHECKING:
+    from .zip import ZippedDir, ZippedFile
+
+
+@dataclass
+class DatasetFile:
+    """Filesystem-backed file entry mirroring zip.ZippedFile API."""
+
+    name: str
+    path: str  # archive-style path (anchor-aware)
+    fs: "DatasetFS"
+
+    def __repr__(self) -> str:
+        try:
+            full = self.fs.root / self.fs._normalize_relpath(self.path)
+            size = full.stat().st_size
+        except Exception:
+            size = "?"
+        return f"DatasetFile(path='{self.path}', size={size})"
+
+    def is_dir(self) -> bool:
+        return False
+
+    def is_file(self) -> bool:
+        return True
+
+    def open(self) -> IO[bytes]:
+        """Open the file for reading in binary mode."""
+        return self.fs.open_binary(self.path)
+
+    def read(self) -> bytes:
+        with self.open() as f:
+            return f.read()
+
+    def buffer(self) -> io.BytesIO:
+        buf = io.BytesIO(self.read())
+        buf.seek(0)
+        return buf
+
+    def isolate(self) -> zipcore.FileBuffer:
+        """Return a FileBuffer wrapping this file's content."""
+        buf = self.buffer()
+        buf.seek(0)
+        return zipcore.FileBuffer(name=self.name, buffer=buf)
+
+    def extract_to(self, dest: Path) -> Path:
+        """Write this file to a directory or file path."""
+        dest_path = Path(dest)
+        if dest_path.is_dir() or str(dest_path).endswith(os.sep):
+            dest_path = dest_path / self.name
+        dest_path.parent.mkdir(parents=True, exist_ok=True)
+        dest_path.write_bytes(self.read())
+        return dest_path
+
+
+@dataclass
+class DatasetDir:
+    """Filesystem-backed directory entry mirroring zip.ZippedDir API."""
+
+    name: str
+    path: str  # archive-style path (anchor-aware)
+    fs: "DatasetFS"
+
+    def __repr__(self) -> str:
+        dirs = [e for e in self.iterdir() if e.is_dir()]
+        files = [e for e in self.iterdir() if e.is_file()]
+        return f"DatasetDir(path='{self.path}', dirs={len(dirs)}, files={len(files)})"
+
+    def is_dir(self) -> bool:
+        return True
+
+    def is_file(self) -> bool:
+        return False
+
+    def listdir(self) -> List[str]:
+        """List immediate children names (dirs first, then files)."""
+        return self.fs.listdir(self.path)
+
+    def iterdir(self) -> Iterable[Union["DatasetDir", "DatasetFile", "ZippedDir", "ZippedFile"]]:
+        """Iterate over children as objects (dirs first, then files)."""
+        yield from self.fs.iterdir(self.path)
+
+
+@dataclass
+class DatasetFS:
+    """Unified view over a dataset rooted in a directory or zipfile.
+
+    Attributes:
+        root: Dataset root (directory or zipfile path).
+        _mode: Backing mode, either "dir" or "zip".
+        _zip: ZipFile handle when `_mode` is "zip", else None.
+        _anchor: Optional top-level directory name inside the archive.
+    """
+    root: Path
+    _mode: Literal["dir", "zip"]
+    _zip: Optional[zipfile.ZipFile]
+    _anchor: str = field(init=False)
+
+    def __post_init__(self) -> None:
+        self._anchor = self._detect_anchor()
+
+    @classmethod
+    def from_path(cls, path: Union[Path, str]) -> "DatasetFS":
+        """Create a DatasetFS from a directory or zip path.
+
+        Args:
+            path: Filesystem path pointing to a dataset root.
+
+        Returns:
+            DatasetFS bound to the given path.
+
+        Raises:
+            ValueError: If the path is neither a directory nor a valid zip file.
+        """
+        path = Path(path)
+        if path.is_dir():
+            return cls(root=path, _mode="dir", _zip=None)
+        if path.is_file() and zipfile.is_zipfile(path):
+            zf = zipcore.load(path)
+            return cls(root=path, _mode="zip", _zip=zf)
+        raise ValueError(f"Invalid dataset root: {path}")
+
+    # -- helpers
+    def _detect_anchor(self) -> str:
+        """Infer the top-level archive directory name.
+
+        Returns:
+            Anchor name when identifiable, else an empty string.
+        """
+        if self._mode == "dir":
+            return self.root.name
+
+        assert self._zip is not None
+        names = [n.strip("/") for n in self._zip.namelist() if n.strip("/")]
+        if not names:
+            return ""
+        first = names[0].split("/")[0]
+        for n in names[1:]:
+            if not n.startswith(first + "/") and n != first:
+                return ""
+        return first
+
+    def _normalize_relpath(self, relpath: str) -> str:
+        """Remove anchor prefix if present.
+
+        Args:
+            relpath: Archive-relative path that may include the anchor.
+
+        Returns:
+            Anchor-stripped relative path.
+        """
+        relpath = relpath.strip("/")
+        if not self._anchor:
+            return relpath
+        if relpath == self._anchor:
+            return ""
+        prefix = f"{self._anchor}/"
+        if relpath.startswith(prefix):
+            return relpath[len(prefix):]
+        return relpath
+
+    def _ensure_anchor(self, relpath: str) -> str:
+        """Add anchor prefix if missing.
+
+        Args:
+            relpath: Archive-relative path without guaranteed anchor.
+
+        Returns:
+            Path guaranteed to include the anchor when one exists.
+        """
+        relpath = relpath.strip("/")
+        if not self._anchor:
+            return relpath
+        if relpath == self._anchor or relpath.startswith(f"{self._anchor}/"):
+            return relpath
+        return f"{self._anchor}/{relpath}" if relpath else self._anchor
+
+    @property
+    def anchor(self) -> str:
+        return self._anchor
+
+    def strip_anchor(self, relpath: str) -> str:
+        """Remove anchor prefix if present."""
+        return self._normalize_relpath(relpath)
+
+    def add_anchor(self, relpath: str) -> str:
+        """Ensure anchor prefix is present."""
+        return self._ensure_anchor(relpath)
+
+    # -- public API
+    def walk(
+        self,
+        top: str = "",
+        *,
+        as_objects: bool = False,
+    ) -> Iterable[Tuple[str, List, List]]:
+        """Yield (dirpath, direntries, fileentries) with archive-style paths.
+
+        Args:
+            top: Optional subdirectory to start from (anchor-aware).
+            as_objects: When True, return DatasetDir/ZippedDir and
+                DatasetFile/ZippedFile entries; otherwise return name strings.
+
+        Yields:
+            Tuples of `(dirpath, direntries, fileentries)` using posix-style paths.
+        """
+        norm_top = top.strip("/")
+        if self._anchor and norm_top and not norm_top.startswith(self._anchor):
+            norm_top = f"{self._anchor}/{norm_top}"
+
+        if self._mode == "dir":
+            base = self.root
+            rel_top = self._normalize_relpath(norm_top)
+            start = base / rel_top if rel_top else base
+            if not start.exists():
+                return
+
+            if not norm_top and self._anchor:
+                # mirror zip.walk: expose the anchor as the top-level directory
+                if as_objects:
+                    yield "", [DatasetDir(name=self._anchor, path=self._anchor, fs=self)], []
+                else:
+                    yield "", [self._anchor], []
+
+            for dirpath, dirnames, filenames in os.walk(start):
+                rel = os.path.relpath(dirpath, base)
+                rel = "" if rel == "." else rel.replace(os.sep, "/")
+                rel = self._ensure_anchor(rel)
+                dirnames = sorted(dirnames)
+                filenames = sorted(filenames)
+                if as_objects:
+                    dir_objs = [
+                        DatasetDir(name=d, path=(f"{rel}/{d}".strip("/")), fs=self) for d in dirnames
+                    ]
+                    file_objs = [
+                        DatasetFile(
+                            name=f, path=(f"{rel}/{f}".strip("/")), fs=self
+                        )
+                        for f in filenames
+                    ]
+                    yield rel, dir_objs, file_objs
+                else:
+                    yield rel, dirnames, filenames
+        else:
+            assert self._zip is not None
+            for dirpath, direntries, files in zipcore.walk(self._zip, top=norm_top):
+                if as_objects:
+                    yield dirpath, direntries, files
+                else:
+                    yield dirpath, [d.name for d in direntries], [f.name for f in files]
+
+    def open_binary(self, relpath: str) -> IO[bytes]:
+        """Open a file by archive-relative path.
+
+        Args:
+            relpath: Path relative to the dataset root (posix separators).
+
+        Returns:
+            File-like object in binary mode.
+
+        Raises:
+            FileNotFoundError: If the file does not exist.
+        """
+        relpath = self._normalize_relpath(relpath)
+
+        if self._mode == "dir":
+            full = self.root / relpath
+            return open(full, "rb")
+        else:
+            assert self._zip is not None
+            arcname = self._ensure_anchor(relpath)
+            top = os.path.dirname(arcname)
+            leaf = os.path.basename(arcname)
+            matches = zipcore.fetch_files_in_zip(
+                self._zip,
+                leaf,
+                top=top,
+                wildcard=False,
+            )
+            if not matches:
+                raise FileNotFoundError(arcname)
+            return matches[0].open()
+
+    def listdir(self, relpath: str = "") -> List[str]:
+        """Return entry names under a relative path (dirs first, then files)."""
+        relpath = self._normalize_relpath(relpath)
+        target = self._ensure_anchor(relpath)
+
+        if self._mode == "dir":
+            base_path = self.root / relpath if relpath else self.root
+            if not base_path.exists():
+                return []
+
+            dirnames: List[str] = []
+            filenames: List[str] = []
+            for entry in base_path.iterdir():
+                if entry.is_dir():
+                    dirnames.append(entry.name)
+                else:
+                    filenames.append(entry.name)
+            dirnames.sort()
+            filenames.sort()
+
+            return dirnames + filenames
+
+        # zip mode
+        assert self._zip is not None
+        dirnames: List[str] = []
+        filenames: List[str] = []
+        dirobjs: List[zipcore.ZippedDir] = []
+        fileobjs: List[zipcore.ZippedFile] = []
+        for dirpath, direntries, files in zipcore.walk(self._zip, top=target):
+            if dirpath != target:
+                continue
+            dirobjs = sorted(direntries, key=lambda d: d.name)
+            fileobjs = sorted(files, key=lambda f: f.name)
+            dirnames = [d.name for d in dirobjs]
+            filenames = [f.name for f in fileobjs]
+            break
+
+        return dirnames + filenames
+
+    def iterdir(
+        self,
+        relpath: str = "",
+    ) -> Iterable[Union["DatasetDir", "DatasetFile", zipcore.ZippedDir, zipcore.ZippedFile]]:
+        """Iterate entries under a relative path as objects (dirs first)."""
+        relpath = self._normalize_relpath(relpath)
+        target = self._ensure_anchor(relpath)
+
+        if self._mode == "dir":
+            base_path = self.root / relpath if relpath else self.root
+            if not base_path.exists():
+                return iter(())
+
+            dir_entries: List[DatasetDir] = []
+            file_entries: List[DatasetFile] = []
+            for entry in base_path.iterdir():
+                name = entry.name
+                p = f"{relpath}/{name}".strip("/")
+                if entry.is_dir():
+                    dir_entries.append(DatasetDir(name=name, path=self._ensure_anchor(p), fs=self))
+                else:
+                    file_entries.append(DatasetFile(name=name, path=self._ensure_anchor(p), fs=self))
+            dir_entries.sort(key=lambda d: d.name)
+            file_entries.sort(key=lambda f: f.name)
+            return iter([*dir_entries, *file_entries])
+
+        # zip mode
+        assert self._zip is not None
+        for dirpath, direntries, files in zipcore.walk(self._zip, top=target):
+            if dirpath != target:
+                continue
+            sorted_dirs = sorted(direntries, key=lambda d: d.name)
+            sorted_files = sorted(files, key=lambda f: f.name)
+            return iter([*sorted_dirs, *sorted_files])
+        return iter(())
+
+    def exists(self, relpath: str) -> bool:
+        """Check existence of a dataset-relative path."""
+        relpath = self._normalize_relpath(relpath)
+        if self._mode == "dir":
+            return (self.root / relpath).exists()
+        else:
+            assert self._zip is not None
+            try:
+                self._zip.getinfo(self._ensure_anchor(relpath))
+                return True
+            except KeyError:
+                return False
+
+    def compress_to(
+        self,
+        dest: Path,
+        *,
+        relpath: str = "",
+        add_root: bool = True,
+        root_name: Optional[str] = None,
+    ) -> Path:
+        """Persist the whole dataset or a subtree as a zip file.
+
+        Args:
+            dest: Destination zip path.
+            relpath: Optional subtree to pack relative to the dataset root.
+            add_root: Whether to include a top-level root folder in the zip.
+            root_name: Optional name for the root folder when `add_root` is True.
+
+        Returns:
+            Path to the created zip file.
+
+        Raises:
+            FileNotFoundError: When the requested subtree does not exist.
+            RuntimeError: When root detection inside a temporary zip fails.
+        """
+        dest = Path(dest)
+        relpath = self._normalize_relpath(relpath)
+        root_name = root_name or self.anchor or (Path(relpath).parts[0] if relpath else self.root.name)
+        relpath = relpath.strip("/")
+
+        with TemporaryDirectory() as tmp:
+            tmp_root = Path(tmp) / root_name
+            if relpath:
+                extract_into = tmp_root / relpath
+            else:
+                extract_into = tmp_root
+
+            if self._mode == "dir":
+                src_dir = self.root / relpath
+                if not src_dir.exists():
+                    raise FileNotFoundError(src_dir)
+                extract_into.parent.mkdir(parents=True, exist_ok=True)
+                shutil.copytree(src_dir, extract_into, dirs_exist_ok=True)
+            else:
+                assert self._zip is not None
+                arcdir = self._ensure_anchor(relpath)
+                if not relpath or arcdir == self.anchor:
+                    # whole zip; just copy
+                    shutil.copyfile(self.root, dest)
+                    return dest
+
+                dirs = zipcore.fetch_dirs_in_zip(
+                    self._zip,
+                    dirname=arcdir,
+                    match_scope="fullpath",
+                    wildcard=False,
+                )
+                target = dirs[0] if dirs else None
+                if target is None:
+                    raise FileNotFoundError(arcdir)
+                with target.isolate() as subzip:
+                    extract_into.parent.mkdir(parents=True, exist_ok=True)
+                    subzip.extractall(extract_into)
+
+            if add_root:
+                tmp_zip = dest.with_suffix(".tmp.zip")
+                zipcore.create_from_dir(tmp_zip, tmp_root)
+                try:
+                    with zipcore.load(tmp_zip) as zf:
+                        roots = zipcore.fetch_dirs_in_zip(
+                            zf, dirname="", match_scope="fullpath", wildcard=True
+                        )
+                        root_dir = roots[0] if roots else None
+                        if root_dir is None:
+                            raise RuntimeError("Failed to locate root dir while zipping.")
+                        root_dir.to_filename(
+                            dest,
+                            add_root=True,
+                            root_name=root_name,
+                            include_dir_entries=True,
+                        )
+                finally:
+                    Path(tmp_zip).unlink(missing_ok=True)
+            else:
+                pack_dir = extract_into if relpath else tmp_root
+                zipcore.create_from_dir(dest, pack_dir)
+
+            return dest
+
+
+__all__ = [
+    "DatasetFS",
+    "DatasetDir",
+    "DatasetFile",
+]
+
+def __dir__() -> List[str]:
+    return sorted(__all__)