glitchlings 1.0.0__cp313-cp313-win_amd64.whl

glitchlings/dlc/pytorch_lightning.py

@@ -0,0 +1,173 @@
+"""Integration helpers for PyTorch Lightning data modules."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any, cast
+
+from ..compat.loaders import get_pytorch_lightning_datamodule
+from ..util.adapters import coerce_gaggle
+from ..zoo import Gaggle, Glitchling
+from ._shared import normalize_column_spec, wrap_dataloader
+
+
+def _glitch_datamodule(
+    datamodule: Any,
+    glitchlings: Glitchling | Gaggle | str | Iterable[str | Glitchling],
+    column: str | Sequence[str],
+    *,
+    seed: int = 151,
+) -> Any:
+    """Return a proxy that applies glitchlings to batches from the datamodule."""
+
+    columns = normalize_column_spec(column)
+    if columns is None:  # pragma: no cover - defensive
+        raise ValueError("At least one column must be specified")
+    # Lightning datamodules only support string column names (mapping keys)
+    columns_str = cast(list[str], columns)
+    gaggle = coerce_gaggle(glitchlings, seed=seed)
+
+    return _GlitchedLightningDataModule(datamodule, columns_str, gaggle)
+
+
+def GlitchedLightningDataModule(
+    datamodule: Any,
+    glitchlings: Glitchling | Gaggle | str | Iterable[str | Glitchling],
+    *,
+    column: str | Sequence[str],
+    seed: int = 151,
+) -> Any:
+    """Return a glitched wrapper around a PyTorch Lightning LightningDataModule.
+
+    This function wraps a LightningDataModule to apply glitchlings to specified
+    columns in batches yielded by the module's dataloaders.
+
+    Args:
+        datamodule: The LightningDataModule to wrap.
+        glitchlings: A glitchling, gaggle, or specification of glitchlings to apply.
+        column: The column name (string) or names (sequence of strings) to corrupt.
+        seed: RNG seed for deterministic corruption (default: 151).
+
+    Returns:
+        A wrapped datamodule that yields corrupted batches from its dataloaders.
+
+    Example:
+        >>> from pytorch_lightning import LightningDataModule
+        >>> from glitchlings.dlc.pytorch_lightning import GlitchedLightningDataModule
+        >>> class MyDataModule(LightningDataModule):
+        ...     def train_dataloader(self):
+        ...         return [{"text": "hello", "label": 0}]
+        >>> dm = MyDataModule()
+        >>> glitched = GlitchedLightningDataModule(dm, "typogre", column="text")
+        >>> batches = list(glitched.train_dataloader())
+    """
+    return _glitch_datamodule(datamodule, glitchlings, column, seed=seed)
+
+
+class _GlitchedLightningDataModule:
+    """Proxy wrapper around a LightningDataModule applying glitchlings to batches."""
+
+    def __init__(self, base: Any, columns: list[str], gaggle: Gaggle) -> None:
+        object.__setattr__(self, "_glitch_base", base)
+        object.__setattr__(self, "_glitch_columns", columns)
+        object.__setattr__(self, "_glitch_gaggle", gaggle)
+
+    def __getattr__(self, attribute: str) -> Any:
+        return getattr(self._glitch_base, attribute)
+
+    def __setattr__(self, attribute: str, value: Any) -> None:
+        if attribute.startswith("_glitch_"):
+            object.__setattr__(self, attribute, value)
+        else:
+            setattr(self._glitch_base, attribute, value)
+
+    def __delattr__(self, attribute: str) -> None:
+        if attribute.startswith("_glitch_"):
+            object.__delattr__(self, attribute)
+        else:
+            delattr(self._glitch_base, attribute)
+
+    def __dir__(self) -> list[str]:
+        return sorted(set(dir(self.__class__)) | set(dir(self._glitch_base)))
+
+    # LightningDataModule API -------------------------------------------------
+    def prepare_data(self, *args: Any, **kwargs: Any) -> Any:
+        return self._glitch_base.prepare_data(*args, **kwargs)
+
+    def setup(self, *args: Any, **kwargs: Any) -> Any:
+        return self._glitch_base.setup(*args, **kwargs)
+
+    def teardown(self, *args: Any, **kwargs: Any) -> Any:
+        return self._glitch_base.teardown(*args, **kwargs)
+
+    def state_dict(self) -> Mapping[str, Any]:
+        state = self._glitch_base.state_dict()
+        return cast(Mapping[str, Any], state)
+
+    def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
+        self._glitch_base.load_state_dict(state_dict)
+
+    def transfer_batch_to_device(self, batch: Any, device: Any, dataloader_idx: int) -> Any:
+        return self._glitch_base.transfer_batch_to_device(batch, device, dataloader_idx)
+
+    def on_before_batch_transfer(self, batch: Any, dataloader_idx: int) -> Any:
+        return self._glitch_base.on_before_batch_transfer(batch, dataloader_idx)
+
+    def on_after_batch_transfer(self, batch: Any, dataloader_idx: int) -> Any:
+        return self._glitch_base.on_after_batch_transfer(batch, dataloader_idx)
+
+    def train_dataloader(self, *args: Any, **kwargs: Any) -> Any:
+        loader = self._glitch_base.train_dataloader(*args, **kwargs)
+        return wrap_dataloader(loader, self._glitch_columns, self._glitch_gaggle)
+
+    def val_dataloader(self, *args: Any, **kwargs: Any) -> Any:
+        loader = self._glitch_base.val_dataloader(*args, **kwargs)
+        return wrap_dataloader(loader, self._glitch_columns, self._glitch_gaggle)
+
+    def test_dataloader(self, *args: Any, **kwargs: Any) -> Any:
+        loader = self._glitch_base.test_dataloader(*args, **kwargs)
+        return wrap_dataloader(loader, self._glitch_columns, self._glitch_gaggle)
+
+    def predict_dataloader(self, *args: Any, **kwargs: Any) -> Any:
+        loader = self._glitch_base.predict_dataloader(*args, **kwargs)
+        return wrap_dataloader(loader, self._glitch_columns, self._glitch_gaggle)
+
+
+# Module initialization: set up inheritance from LightningDataModule if available
+def _setup_inheritance() -> None:
+    """Set up _GlitchedLightningDataModule to inherit from LightningDataModule.
+
+    This function is called once at module import time to dynamically set the base
+    class of _GlitchedLightningDataModule to inherit from
+    pytorch_lightning.LightningDataModule when available. This ensures that
+    isinstance(glitched, LightningDataModule) checks work correctly and that the
+    wrapper interoperates with Lightning APIs that require that type.
+    """
+    datamodule_cls = get_pytorch_lightning_datamodule()
+    if datamodule_cls is None:
+        # If LightningDataModule is not available, keep as plain object
+        return
+
+    # Try to dynamically set __bases__ to inherit from LightningDataModule
+    try:
+        _GlitchedLightningDataModule.__bases__ = (datamodule_cls,)
+    except TypeError:
+        # If we can't modify __bases__ (e.g., due to __slots__), create a new class
+        namespace = {
+            name: value
+            for name, value in vars(_GlitchedLightningDataModule).items()
+            if name not in {"__dict__", "__weakref__"}
+        }
+        replacement = cast(
+            type[Any],
+            type("_GlitchedLightningDataModule", (datamodule_cls,), namespace),
+        )
+        # Update the module's global namespace
+        globals()["_GlitchedLightningDataModule"] = replacement
+
+
+# Set up inheritance at module import time
+_setup_inheritance()
+
+
+__all__ = ["GlitchedLightningDataModule"]

glitchlings/internal/__init__.py

@@ -0,0 +1,16 @@
+"""Internal utilities shared across the glitchlings package.
+
+This subpackage contains impure modules that handle side effects:
+
+- ``rust.py``: Low-level Rust extension loader and FFI primitives
+- ``rust_ffi.py``: High-level Rust operation wrappers (preferred entry point)
+
+Pure modules should NOT import from this package. Use the operations in
+``rust_ffi.py`` at boundary layers only.
+
+See AGENTS.md "Functional Purity Architecture" for details.
+"""
+
+from __future__ import annotations
+
+__all__ = []

glitchlings/internal/rust.py

@@ -0,0 +1,159 @@
+"""Shared helpers for loading the compiled Rust extension."""
+
+from __future__ import annotations
+
+import random
+import sys
+from importlib import machinery, util
+from pathlib import Path
+from types import ModuleType
+from typing import Any, Callable, Mapping, MutableMapping, cast
+
+_EXTENSION_STEM = "_corruption_engine"
+
+
+class RustExtensionImportError(RuntimeError):
+    """Raised when the compiled Rust extension cannot be imported."""
+
+
+def _iter_extension_candidates() -> tuple[Path, ...]:
+    """Return likely paths for the compiled extension within the package."""
+
+    package_root = Path(__file__).resolve().parents[1]
+    extension_dir = package_root / _EXTENSION_STEM
+    search_roots = (extension_dir, package_root)
+
+    candidates: list[Path] = []
+    for root in search_roots:
+        for suffix in machinery.EXTENSION_SUFFIXES:
+            candidate = (root / _EXTENSION_STEM).with_suffix(suffix)
+            if candidate.exists():
+                candidates.append(candidate)
+    return tuple(candidates)
+
+
+def _existing_compiled_module() -> ModuleType | None:
+    """Return a previously loaded compiled module if one is present."""
+
+    for name in ("glitchlings._corruption_engine", "_corruption_engine"):
+        module = sys.modules.get(name)
+        if module is None:
+            continue
+        module_file = getattr(module, "__file__", "")
+        if module_file and not str(module_file).endswith("__init__.py"):
+            return module
+    return None
+
+
+def _load_extension_from_disk() -> ModuleType:
+    """Load the compiled extension from disk or raise if unavailable."""
+
+    candidates = _iter_extension_candidates()
+    for candidate in candidates:
+        spec = util.spec_from_file_location(_EXTENSION_STEM, candidate)
+        if spec is None or spec.loader is None:
+            continue
+        module = util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        return module
+
+    searched = ", ".join(str(path) for path in candidates) or "<unavailable>"
+    message = (
+        "glitchlings._corruption_engine failed to import. Rebuild the project with"
+        "`pip install .` or `maturin develop` so the compiled extension is available "
+        f"(searched: {searched})."
+    )
+    raise RustExtensionImportError(message)
+
+
+def load_rust_module() -> ModuleType:
+    """Return the compiled Rust module, loading it on demand."""
+
+    existing = _existing_compiled_module()
+    if existing is not None:
+        return existing
+
+    module = _load_extension_from_disk()
+    sys.modules.setdefault("glitchlings._corruption_engine", module)
+    sys.modules.setdefault("_corruption_engine", module)
+    return module
+
+
+_RUST_MODULE: ModuleType | None = None
+_OPERATION_CACHE: MutableMapping[str, Callable[..., Any]] = {}
+
+
+def _get_rust_module() -> ModuleType:
+    """Return the compiled Rust module, importing it on first use."""
+
+    global _RUST_MODULE
+
+    if _RUST_MODULE is None:
+        _RUST_MODULE = load_rust_module()
+
+    return _RUST_MODULE
+
+
+def _build_missing_operation_error(name: str) -> RuntimeError:
+    message = (
+        "Rust operation '{name}' is not exported by glitchlings._corruption_engine."
+        "Rebuild the project to refresh the compiled extension."
+    )
+    return RuntimeError(message.format(name=name))
+
+
+def resolve_seed(seed: int | None, rng: random.Random | None) -> int:
+    """Resolve a 64-bit seed using an optional RNG."""
+
+    if seed is not None:
+        return int(seed) & 0xFFFFFFFFFFFFFFFF
+    if rng is not None:
+        return rng.getrandbits(64)
+    return random.getrandbits(64)
+
+
+def get_rust_operation(operation_name: str) -> Callable[..., Any]:
+    """Return a callable exported by :mod:`glitchlings._corruption_engine`.
+
+    Parameters
+    ----------
+    operation_name : str
+        Name of the function to retrieve from the compiled extension.
+
+    Raises
+    ------
+    RuntimeError
+        If the operation cannot be located or is not callable.
+    """
+
+    operation = _OPERATION_CACHE.get(operation_name)
+    if operation is not None:
+        return operation
+
+    module = _get_rust_module()
+    try:
+        candidate = getattr(module, operation_name)
+    except AttributeError as exc:
+        raise _build_missing_operation_error(operation_name) from exc
+
+    if not callable(candidate):
+        raise _build_missing_operation_error(operation_name)
+
+    operation = cast(Callable[..., Any], candidate)
+    _OPERATION_CACHE[operation_name] = operation
+    return operation
+
+
+def preload_operations(*operation_names: str) -> Mapping[str, Callable[..., Any]]:
+    """Eagerly load multiple Rust operations at once."""
+
+    return {name: get_rust_operation(name) for name in operation_names}
+
+
+__all__ = [
+    "RustExtensionImportError",
+    "get_rust_operation",
+    "load_rust_module",
+    "preload_operations",
+    "resolve_seed",
+]