weftcodec 1.2.0__py3-none-any.whl

weftcodec/torch.py

@@ -0,0 +1,94 @@
+"""PyTorch-facing dataset over a mounted WeftFS directory.
+
+``torch`` is an optional dependency and is never imported at module import time.
+Plain ``bytes`` access works without it; requesting a tensor raises a clear
+``ImportError`` when PyTorch is absent.
+"""
+
+from __future__ import annotations
+
+import importlib
+import os
+from pathlib import Path
+
+
+def _torch_module():
+    try:
+        return importlib.import_module("torch")
+    except ImportError as exc:  # pragma: no cover - exercised without torch
+        raise ImportError(
+            "weftcodec.torch requires PyTorch; install it with "
+            "`pip install weftcodec[torch]` (or `pip install torch`)"
+        ) from exc
+
+
+class WeftMapDataset:
+    """Deterministic map-style dataset over files in a WeftFS mount.
+
+    Parameters
+    ----------
+    root:
+        Directory served by a WeftFS mount (or any directory).
+    transform:
+        Optional callable applied to the raw ``bytes`` before the item is
+        returned or converted to a tensor.
+    dtype:
+        When set (for example ``torch.uint8`` or ``torch.float32``), items are
+        returned as ``torch.frombuffer`` tensors instead of ``bytes``.
+    shape:
+        Optional shape used to reshape a returned tensor. ``-1`` can be used
+        for one inferred dimension.
+    """
+
+    def __init__(
+        self,
+        root: "str | os.PathLike[str]",
+        transform=None,
+        *,
+        dtype=None,
+        shape: "tuple[int, ...] | None" = None,
+    ) -> None:
+        self.root = Path(root)
+        if not self.root.is_dir():
+            raise NotADirectoryError("dataset root is not a directory: %s" % self.root)
+        self.transform = transform
+        self.dtype = dtype
+        self.shape = tuple(shape) if shape is not None else None
+        self._entries = self._walk()
+
+    def _walk(self):
+        entries = []
+        for dirpath, dirnames, filenames in os.walk(self.root):
+            dirnames.sort()
+            for name in sorted(filenames):
+                full = Path(dirpath) / name
+                if full.is_file():
+                    relative = full.relative_to(self.root).as_posix()
+                    entries.append((relative, full))
+        entries.sort(key=lambda entry: entry[0])
+        return entries
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+    def paths(self) -> "list[str]":
+        return [relative for relative, _ in self._entries]
+
+    def __getitem__(self, index: int):
+        if index < 0:
+            index += len(self._entries)
+        if index < 0 or index >= len(self._entries):
+            raise IndexError("dataset index %d out of range" % index)
+        relative, full = self._entries[index]
+        with open(full, "rb") as handle:
+            data = handle.read()
+        if self.transform is not None:
+            data = self.transform(data)
+        if self.dtype is None:
+            return relative, data
+        torch = _torch_module()
+        # `bytes` is read-only; torch.frombuffer needs a writable buffer.
+        tensor = torch.frombuffer(bytearray(data), dtype=self.dtype)
+        if self.shape is not None:
+            tensor = tensor.reshape(self.shape)
+        return relative, tensor

weftcodec/varint.py

@@ -0,0 +1,88 @@
+"""Minimal unsigned LEB128 varints (WEFTCODEC-FORMAT §1).
+
+The decoder mirrors the Rust implementation exactly:
+
+* at most 10 bytes;
+* a multi-byte encoding must not end in ``0x00`` (``NonCanonical``);
+* the tenth byte must be ``0x00`` or ``0x01`` (``Malformed`` otherwise);
+* truncated input is ``Malformed``.
+"""
+
+from __future__ import annotations
+
+from .errors import Malformed, NonCanonical
+
+MAX_VARINT_LEN = 10
+U64_MAX = (1 << 64) - 1
+
+
+def encoded_len(value: int) -> int:
+    if value < 0 or value > U64_MAX:
+        raise Malformed("varint value is outside u64")
+    n = 1
+    while value >= 0x80:
+        value >>= 7
+        n += 1
+    return n
+
+
+def encode(value: int, out: bytearray) -> None:
+    """Append the minimal LEB128 encoding of ``value`` to ``out``."""
+    if value < 0 or value > U64_MAX:
+        raise Malformed("varint value is outside u64")
+    v = value
+    while True:
+        byte = v & 0x7F
+        v >>= 7
+        if v == 0:
+            out.append(byte)
+            return
+        out.append(byte | 0x80)
+
+
+def encode_bytes(value: int) -> bytes:
+    out = bytearray()
+    encode(value, out)
+    return bytes(out)
+
+
+def decode(data: bytes, pos: int) -> "tuple[int, int]":
+    """Decode one varint at ``pos``; return ``(value, next_pos)``."""
+    value = 0
+    shift = 0
+    offset = pos
+    count = 0
+    while True:
+        if offset >= len(data):
+            raise Malformed("truncated varint")
+        byte = data[offset]
+        offset += 1
+        count += 1
+        if count > MAX_VARINT_LEN:
+            raise Malformed("varint longer than 10 bytes")
+        if count == MAX_VARINT_LEN and byte > 0x01:
+            raise Malformed("varint overflows u64")
+        low = byte & 0x7F
+        if shift < 64:
+            value |= low << shift
+        elif low != 0:
+            raise Malformed("varint overflows u64")
+        if byte & 0x80 == 0:
+            if count > 1 and byte == 0x00:
+                raise NonCanonical("non-minimal varint")
+            return value, offset
+        shift += 7
+        if shift > 63:
+            raise Malformed("varint exceeds 10 bytes")
+
+
+def zigzag_encode(value: int) -> int:
+    """``(v << 1) ^ (v >> 63)`` on a 64-bit two's-complement pattern."""
+    if value < -(1 << 63) or value > (1 << 63) - 1:
+        raise Malformed("signed value is outside i64")
+    return ((value << 1) ^ (value >> 63)) & U64_MAX
+
+
+def zigzag_decode(value: int) -> int:
+    """Inverse of :func:`zigzag_encode`."""
+    return (value >> 1) ^ -(value & 1)

weftcodec/wcb.py

@@ -0,0 +1,310 @@
+"""Canonical WCB bytes (WEFTCODEC-FORMAT §8).
+
+WCB is a strict, explicit subset of deterministic CBOR. Python values map as:
+
+===============  ==========================================
+Python           WCB
+===============  ==========================================
+``None``         null (tag 0x00)
+``False``        false (0x01)
+``True``         true (0x02)
+``U64(n)``       unsigned int (0x03)
+``I64(n)``       signed int (0x04)
+``str``          text (0x05), UTF-8
+``bytes``        bytes (0x06)
+``list``/tuple   array (0x07)
+``dict``         map (0x08), text keys only
+``Timestamp(n)`` timestamp (0x09), i64 nanoseconds since the epoch
+===============  ==========================================
+
+Tag ``0x0A..0xFF`` is ``NonCanonical`` in v1. Map keys are sorted by their
+full encoded bytes, duplicates are rejected, and varints are minimal.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from . import varint
+from .errors import LimitExceeded, Malformed, NonCanonical
+from .limits import Limits
+
+TAG_NULL = 0x00
+TAG_FALSE = 0x01
+TAG_TRUE = 0x02
+TAG_U64 = 0x03
+TAG_I64 = 0x04
+TAG_TEXT = 0x05
+TAG_BYTES = 0x06
+TAG_ARRAY = 0x07
+TAG_MAP = 0x08
+TAG_TIMESTAMP = 0x09
+
+I64_MIN = -(1 << 63)
+I64_MAX = (1 << 63) - 1
+U64_MAX = (1 << 64) - 1
+
+
+@dataclass(frozen=True)
+class U64:
+    """Unsigned 64-bit integer."""
+
+    value: int
+
+
+@dataclass(frozen=True)
+class I64:
+    """Signed 64-bit integer."""
+
+    value: int
+
+
+@dataclass(frozen=True)
+class Timestamp:
+    """Signed 64-bit nanoseconds since the Unix epoch, UTC."""
+
+    nanos: int
+
+
+def encode(value, limits: "Limits | None" = None) -> bytes:
+    """Encode ``value`` to canonical WCB bytes."""
+    encoder = _Encoder(limits or Limits())
+    out = bytearray()
+    encoder.value(value, 0, out)
+    return bytes(out)
+
+
+def decode(data: bytes, limits: "Limits | None" = None):
+    """Decode canonical WCB bytes, rejecting every non-canonical form."""
+    if not data:
+        raise Malformed("empty WCB value")
+    decoder = _Decoder(data, limits or Limits())
+    value = decoder.value(0)
+    if decoder.pos != len(data):
+        raise NonCanonical(
+            "trailing bytes after WCB value: %d of %d unconsumed"
+            % (len(data) - decoder.pos, len(data))
+        )
+    return value
+
+
+class _Encoder:
+    __slots__ = ("limits", "items")
+
+    def __init__(self, limits: Limits) -> None:
+        self.limits = limits
+        self.items = 0
+
+    def _bump_item(self) -> None:
+        self.items += 1
+        if self.items > self.limits.max_structured_items:
+            raise LimitExceeded(
+                "WCB item count %d exceeds limit %d"
+                % (self.items, self.limits.max_structured_items)
+            )
+
+    def _check_text(self, length: int, what: str) -> None:
+        if length > self.limits.max_string_bytes:
+            raise LimitExceeded(
+                "%s length %d exceeds limit %d"
+                % (what, length, self.limits.max_string_bytes)
+            )
+
+    def _enter(self, depth: int) -> int:
+        nxt = depth + 1
+        if nxt > self.limits.max_structured_depth:
+            raise LimitExceeded(
+                "WCB nesting depth %d exceeds limit %d"
+                % (nxt, self.limits.max_structured_depth)
+            )
+        return nxt
+
+    def value(self, value, depth: int, out: bytearray) -> None:
+        self._bump_item()
+        if value is None:
+            out.append(TAG_NULL)
+        elif value is False:
+            out.append(TAG_FALSE)
+        elif value is True:
+            out.append(TAG_TRUE)
+        elif isinstance(value, U64):
+            if value.value < 0 or value.value > U64_MAX:
+                raise Malformed("U64 value is outside u64")
+            out.append(TAG_U64)
+            varint.encode(value.value, out)
+        elif isinstance(value, I64):
+            if value.value < I64_MIN or value.value > I64_MAX:
+                raise Malformed("I64 value is outside i64")
+            out.append(TAG_I64)
+            varint.encode(varint.zigzag_encode(value.value), out)
+        elif isinstance(value, Timestamp):
+            if value.nanos < I64_MIN or value.nanos > I64_MAX:
+                raise Malformed("Timestamp value is outside i64")
+            out.append(TAG_TIMESTAMP)
+            varint.encode(varint.zigzag_encode(value.nanos), out)
+        elif isinstance(value, str):
+            raw = value.encode("utf-8")
+            self._check_text(len(raw), "text")
+            out.append(TAG_TEXT)
+            varint.encode(len(raw), out)
+            out.extend(raw)
+        elif isinstance(value, (bytes, bytearray, memoryview)):
+            raw = bytes(value)
+            self._check_text(len(raw), "bytes")
+            out.append(TAG_BYTES)
+            varint.encode(len(raw), out)
+            out.extend(raw)
+        elif isinstance(value, (list, tuple)):
+            child = self._enter(depth)
+            if len(value) > self.limits.max_structured_items:
+                raise LimitExceeded("WCB array too large")
+            out.append(TAG_ARRAY)
+            varint.encode(len(value), out)
+            for item in value:
+                self.value(item, child, out)
+        elif isinstance(value, dict):
+            child = self._enter(depth)
+            if len(value) > self.limits.max_structured_items:
+                raise LimitExceeded("WCB map too large")
+            entries = []
+            for key in value:
+                if not isinstance(key, str):
+                    raise Malformed("WCB map keys must be text")
+                raw = key.encode("utf-8")
+                self._check_text(len(raw), "map key")
+                encoded_key = bytearray()
+                encoded_key.append(TAG_TEXT)
+                varint.encode(len(raw), encoded_key)
+                encoded_key.extend(raw)
+                entries.append((bytes(encoded_key), key))
+            entries.sort(key=lambda entry: entry[0])
+            out.append(TAG_MAP)
+            varint.encode(len(entries), out)
+            for encoded_key, key in entries:
+                out.extend(encoded_key)
+                self.value(value[key], child, out)
+        else:
+            raise Malformed("value of type %s is not representable in WCB" % type(value).__name__)
+
+
+class _Decoder:
+    __slots__ = ("input", "pos", "limits", "items")
+
+    def __init__(self, data: bytes, limits: Limits) -> None:
+        self.input = data
+        self.pos = 0
+        self.limits = limits
+        self.items = 0
+
+    def _bump_item(self) -> None:
+        self.items += 1
+        if self.items > self.limits.max_structured_items:
+            raise LimitExceeded(
+                "WCB item count %d exceeds limit %d"
+                % (self.items, self.limits.max_structured_items)
+            )
+
+    def _read_u8(self) -> int:
+        if self.pos >= len(self.input):
+            raise Malformed("truncated WCB value")
+        byte = self.input[self.pos]
+        self.pos += 1
+        return byte
+
+    def _read_varint(self) -> int:
+        value, self.pos = varint.decode(self.input, self.pos)
+        return value
+
+    def _remaining(self) -> int:
+        return len(self.input) - self.pos
+
+    def _take(self, length: int, what: str) -> bytes:
+        if length > self.limits.max_string_bytes:
+            raise LimitExceeded(
+                "%s length %d exceeds limit %d"
+                % (what, length, self.limits.max_string_bytes)
+            )
+        if length > self._remaining():
+            raise Malformed("truncated WCB %s" % what)
+        start = self.pos
+        self.pos += length
+        return self.input[start : self.pos]
+
+    def _enter(self, depth: int) -> int:
+        nxt = depth + 1
+        if nxt > self.limits.max_structured_depth:
+            raise LimitExceeded(
+                "WCB nesting depth %d exceeds limit %d"
+                % (nxt, self.limits.max_structured_depth)
+            )
+        return nxt
+
+    def _check_count(self, count: int, min_bytes_each: int, what: str) -> int:
+        if count > self.limits.max_structured_items:
+            raise LimitExceeded(
+                "WCB %s count %d exceeds limit %d"
+                % (what, count, self.limits.max_structured_items)
+            )
+        if count * min_bytes_each > self._remaining():
+            raise Malformed(
+                "truncated WCB %s: count %d cannot fit in %d remaining bytes"
+                % (what, count, self._remaining())
+            )
+        return count
+
+    def value(self, depth: int):
+        self._bump_item()
+        tag = self._read_u8()
+        if tag == TAG_NULL:
+            return None
+        if tag == TAG_FALSE:
+            return False
+        if tag == TAG_TRUE:
+            return True
+        if tag == TAG_U64:
+            return U64(self._read_varint())
+        if tag == TAG_I64:
+            return I64(varint.zigzag_decode(self._read_varint()))
+        if tag == TAG_TIMESTAMP:
+            return Timestamp(varint.zigzag_decode(self._read_varint()))
+        if tag == TAG_TEXT:
+            raw = self._take(self._read_varint(), "text")
+            try:
+                return raw.decode("utf-8")
+            except UnicodeDecodeError:
+                raise NonCanonical("WCB text is not valid UTF-8")
+        if tag == TAG_BYTES:
+            return self._take(self._read_varint(), "bytes")
+        if tag == TAG_ARRAY:
+            child = self._enter(depth)
+            count = self._check_count(self._read_varint(), 1, "array")
+            return [self.value(child) for _ in range(count)]
+        if tag == TAG_MAP:
+            child = self._enter(depth)
+            count = self._check_count(self._read_varint(), 2, "map")
+            result = {}
+            previous_key = None
+            for _ in range(count):
+                key_start = self.pos
+                key_tag = self._read_u8()
+                if key_tag != TAG_TEXT:
+                    raise NonCanonical("WCB map key is not a text value (tag 0x05)")
+                raw = self._take(self._read_varint(), "map key")
+                key_end = self.pos
+                if previous_key is not None:
+                    previous = self.input[previous_key[0] : previous_key[1]]
+                    current = self.input[key_start:key_end]
+                    if current <= previous:
+                        raise NonCanonical(
+                            "WCB map keys are duplicated or not sorted by encoded key bytes"
+                        )
+                previous_key = (key_start, key_end)
+                try:
+                    key = raw.decode("utf-8")
+                except UnicodeDecodeError:
+                    raise NonCanonical("WCB map key is not valid UTF-8")
+                if key in result:
+                    raise NonCanonical("duplicate WCB map key")
+                result[key] = self.value(child)
+            return result
+        raise NonCanonical("unknown WCB tag 0x%02x" % tag)