kups.core.lens ¶

class BaseLens(Lens[S, R], abc.ABC):
    """Base class for lens implementations."""

    def focus[B](self, where: Callable[[R], B]) -> Lens[S, B]:
        return NestedLens(self, SimpleLens(view(where)))

    def apply[S2](self: BaseLens[S2, R], state: S2, modifier: Modifier[R]) -> S2:
        """Apply a modifier function to the focused value."""
        return self.set(state, modifier(self.get(state)))

    def bind[S2](self: BaseLens[S2, R], state: S2) -> BoundLens[S2, R]:
        """Bind this lens to a specific data structure."""
        return SimpleBoundLens(self, state)

    def at(self, idxs, *, args: ScatterArgs | None = None) -> Lens[S, R]:
        """Create a lens that slices the focused pytree at the given indices."""
        return IndexLens(self, idxs, args=args if args is not None else {})

    def merge[S2, R2](
        self: BaseLens[S2, R], other: Lens[S2, R2]
    ) -> Lens[S2, tuple[R, R2]]:
        return MergedLens(self, other)

    def nest[U, R2](self: Lens[S, R2], other: Lens[R2, U]) -> Lens[S, U]:
        view_func = other.get if isinstance(other, Lens) else other
        return self.focus(view_func)

    @overload
    def __call__[S2, R2](self: Lens[S2, R2], state: S2, /, value: R2) -> S2: ...

    @overload
    def __call__[S2, R2](
        self: Lens[S2, R2], state: S2, /, value: _NOT_SET_TYPE = _NOT_SET
    ) -> R2: ...

    def __call__[S2, R2](
        self: Lens[S2, R2], state: S2, /, value: R2 | _NOT_SET_TYPE = _NOT_SET
    ) -> S2 | R2:
        if value is _NOT_SET:
            return self.get(state)
        else:
            return self.set(state, value)

class BoundLens(Protocol[S, R]):
    """Protocol for a lens that has been bound to a specific data structure.

    A bound lens provides the same operations as a regular lens but without
    requiring the state parameter, since it's already bound to a specific instance.

    Bound lenses satisfy View (via zero-argument call) and can update via single-argument
    call, providing a convenient interface for repeated operations on the same state.

    Generic in ``S`` (bound data structure) and ``R`` (focused value).

    Examples:
        >>> bound = my_lens.bind(state)
        >>> bound()  # View: returns focused value
        >>> bound(new_value)  # Update: returns modified state
    """

    @overload
    def __call__[R2](self: BoundLens[S, R2], value: R2) -> S: ...

    @overload
    def __call__(self) -> R: ...

    def __call__[R2](
        self: BoundLens[S, R2], value: R2 | _NOT_SET_TYPE = _NOT_SET
    ) -> S | R2:
        """Get or set the focused value in the bound state.

        When called with no arguments, returns the focused value.
        When called with one argument, sets the value and returns the modified state.

        Args:
            value: If provided, the new value to set

        Returns:
            The focused value (no args) or modified state (one arg)
        """
        ...

    def focus[B, R2](
        self: BoundLens[S, R2], where: Callable[[R2], B]
    ) -> BoundLens[S, B]:
        """Focus this bound lens on a deeper part of the data structure.

        Args:
            where: A function that extracts a value from the current focus

        Returns:
            A new bound lens focused on the result of the where function
        """
        ...

    def get(self) -> R:
        """Extract the focused value from the bound data structure.

        Returns:
            The focused value from the bound state
        """
        ...

    def set[R2](self: BoundLens[S, R2], value: R2) -> S:
        """Set the focused value in the bound data structure.

        Args:
            value: The new value to set

        Returns:
            A new data structure with the value set
        """
        ...

    def apply[R2](self: BoundLens[S, R2], modifier: Modifier[R2]) -> S:
        """Apply a modifier function to the focused value in the bound data structure.

        Args:
            modifier: A function that transforms the focused value

        Returns:
            A new data structure with the modified value set
        """
        ...

    def at(self, idxs: Any, *, args: ScatterArgs | None = None) -> BoundLens[S, R]:
        """Create a bound lens that slices the focused pytree at the given indices.

        Functionally equivalent to composing jax.tree.map and jax.array.at[idxs].get/set
        on the bound data structure.

        Args:
            idxs: The indices to slice the pytree at.
            args: Optional ScatterArgs controlling scatter behavior (mode,
                wrap_negative_indices, fill_value, indices_are_sorted, unique_indices).

        Returns:
            A new bound lens that slices the focused pytree at the given indices.
        """
        ...

    def merge[S2, R2, R3](
        self: BoundLens[S2, R2], other: Lens[S2, R3]
    ) -> BoundLens[S2, tuple[R2, R3]]:
        """Merge this lens with another lens to access multiple values.

        Args:
            other: Another lens to merge with

        Returns:
            A new lens that accesses both focused values as a tuple
        """
        ...

    def nest[U, R2](self: BoundLens[S, R2], other: Lens[R2, U]) -> BoundLens[S, U]:
        """Nest another lens or view within this bound lens.

        This provides an alternative to focus() that works with both lenses and views.
        When given a lens, it extracts the view from it; when given a view directly,
        it uses it as-is.

        Args:
            other: Either a lens or view to nest within this lens

        Returns:
            A new bound lens that composes this lens with the provided lens/view
        """
        ...

@dataclass
class ConstLens(BaseLens[S, R]):
    """Lens that always returns a constant value; set is a no-op."""

    value: R

    def get[S2](self: ConstLens[S2, R], state: S2) -> R:
        return self.value

    def set[S2, R2](self: ConstLens[S2, R2], state: S2, value: R2) -> S2:
        return state

class HasLensFields(metaclass=FieldMetaAccess):
    """Base class for dataclasses that support lens-enabled field access.

    Dataclasses inheriting from HasLensFields can use LensField[T] annotations
    to enable dual-mode field access:
    - Class access returns a Lens object for functional operations
    - Instance access returns the field value normally

    This class uses the FieldMetaAccess metaclass to intercept class attribute
    access and provide lens objects when appropriate.

    Example:
        >>> from kups.core.utils.jax import dataclass
        >>> from jax import Array
        >>>
        >>> @dataclass
        ... class State(HasLensFields):
        ...     position: LensField[Array]
        ...     velocity: LensField[Array]
        >>>
        >>> state = State(position=pos, velocity=vel)
        >>> pos_lens = State.position  # Returns Lens[State, Array]
        >>> current_pos = state.position  # Returns the Array value

    Note:
        HasLensFields itself cannot be instantiated. It must be subclassed.
    """

    def __new__(cls, *args, **kwargs):
        if cls is HasLensFields:
            raise TypeError(
                "Can't instantiate abstract class {}".format(HasLensFields.__name__)
            )
        return super(HasLensFields, cls).__new__(cls)

@dataclass
class IndexLens(BaseLens[S, R]):
    """A lens that performs array indexing operations on the focused data.

    This lens wraps another lens and applies JAX array indexing operations
    to slice, index, or select specific elements from arrays in the focused data.
    """

    lens: Lens[S, R] = field(static=True)
    idxs: Array
    args: ScatterArgs = field(default_factory=lambda: {}, static=True)

    def focus[B](self, where: Callable[[R], B]) -> Lens[S, B]:
        """Focus deeper on the indexed data."""
        raise RuntimeError(
            "IndexLens cannot be focused further. Please reorder your lenses."
        )

    def get[S2](self: IndexLens[S2, R], state: S2) -> R:
        """Get values by applying array indexing to the focused data."""

        def _array_getter(scatter_args: ScatterArgs, arr: Array):
            return arr.at[self.idxs].get(**scatter_args)

        def _getter(arr: Array | HasScatterArgs):
            if isinstance(arr, Array):
                return _array_getter(self.args, arr)
            return tree_map(
                partial(_array_getter, {**arr.scatter_args, **self.args}), arr
            )

        return tree_map(
            _getter,
            self.lens.get(state),
            is_leaf=lambda x: isinstance(x, (Array, HasScatterArgs)),
        )

    def set[S2, R2](self: IndexLens[S2, R2], state: S2, value: R2) -> S2:
        """Set values by applying array indexing to the focused data."""
        return self.lens.set(
            state, tree_scatter_set(self.lens.get(state), value, self.idxs, self.args)
        )

    def at(self, idxs, **extra_kwargs) -> Lens[S, R]:
        """Create a nested index lens for further slicing."""
        raise RuntimeError(
            "IndexLens cannot be sliced further. Please merge your lenses."
        )

@dataclass
class LambdaLens(BaseLens[S, R]):
    """A lens that uses custom getter and setter functions.

    This allows for more complex lens behavior that cannot be expressed
    with simple field access or traversal-based operations.
    """

    _get: View[S, R] = field(static=True)
    _set: Update[S, R] = field(static=True)

    def get[S2](self: LambdaLens[S2, R], state: S2) -> R:
        """Get the focused value using the custom getter function."""
        return self._get(state)

    def set[S2, R2](self: LambdaLens[S2, R2], state: S2, value: R2) -> S2:
        """Set the focused value using the custom setter function."""
        return self._set(state, value)

@runtime_checkable
class Lens(Protocol[S, R]):
    """Protocol for a lens that provides bidirectional access to data structures.

    A lens combines a getter and setter, allowing functional access and modification
    of nested data structures. Lenses are composable and can be focused on specific
    parts of a data structure.

    Lenses satisfy both the View and Update protocols through their `__call__` method,
    allowing them to be used directly wherever a view or update function is expected.

    Generic in ``S`` (source data structure) and ``R`` (focused value).

    Examples:
        >>> state = MyState(value=10)
        >>> my_lens = lens(lambda s: s.value)
        >>> my_lens(state)  # View: returns 10
        >>> my_lens(state, 20)  # Update: returns MyState(value=20)
    """

    @overload
    def __call__[S2, R2](self: Lens[S2, R2], state: S2, /, value: R2) -> S2: ...

    @overload
    def __call__[S2, R2](self: Lens[S2, R2], state: S2, /) -> R2: ...

    def __call__[S2, R2](
        self: Lens[S2, R2], state: S2, /, value: R2 | _NOT_SET_TYPE = _NOT_SET
    ) -> S2 | R2:
        """Get or set the focused value, satisfying View and Update protocols.

        When called with one argument, acts as a View and returns the focused value.
        When called with two arguments, acts as an Update and returns the modified state.

        Args:
            state: The data structure to operate on
            value: If provided, the new value to set

        Returns:
            The focused value (one arg) or modified state (two args)
        """
        ...

    def focus[B](self, where: Callable[[R], B]) -> Lens[S, B]:
        """Focus this lens on a deeper part of the data structure.

        Args:
            where: A function that extracts a value from the current focus

        Returns:
            A new lens focused on the result of the where function
        """
        ...

    def get[S2](self: Lens[S2, R], state: S2, /) -> R:
        """Extract the focused value from the data structure.

        Args:
            state: The data structure to extract from

        Returns:
            The focused value
        """
        ...

    def set[S2, R2](self: Lens[S2, R2], state: S2, /, value: R2) -> S2:
        """Set the focused value in the data structure.

        Args:
            state: The data structure to modify
            value: The new value to set

        Returns:
            A new data structure with the value set
        """
        ...

    def apply[S2, R2](self: Lens[S2, R2], state: S2, /, modifier: Modifier[R2]) -> S2:
        """Apply a modifier function to the focused value.

        Args:
            state: The data structure to modify
            modifier: A function that transforms the focused value

        Returns:
            A new data structure with the modified value
        """
        ...

    def bind[S2](self: Lens[S2, R], state: S2, /) -> BoundLens[S2, R]:
        """Bind this lens to a specific data structure.

        Args:
            state: The data structure to bind to

        Returns:
            A bound lens that operates on the given state
        """
        ...

    def at(self, idxs: Any, *, args: ScatterArgs | None = None) -> Lens[S, R]:
        """Create a lens that slices the focused pytree at the given indices.

        Functionally equivalent to composing jax.tree.map and jax.array.at[idxs].get/set.

        Args:
            idxs: The indices to slice the pytree at.
            args: Optional ScatterArgs controlling scatter behavior (mode,
                wrap_negative_indices, fill_value, indices_are_sorted, unique_indices).

        Returns:
            A new lens that slices the focused pytree at the given indices.
        """
        ...

    def merge[S2, R2](self: Lens[S2, R], other: Lens[S2, R2]) -> Lens[S2, tuple[R, R2]]:
        """Merge this lens with another lens to access multiple values.

        Args:
            other: Another lens to merge with

        Returns:
            A new lens that accesses both focused values as a tuple
        """
        ...

    def nest[U, S2, R2](self: Lens[S2, R2], other: Lens[R2, U]) -> Lens[S2, U]:
        """Nest another lens or view within this lens.

        This provides an alternative to focus() that works with both lenses and views.
        When given a lens, it extracts the view from it; when given a view directly,
        it uses it as-is.

        Args:
            other: Either a lens or view to nest within this lens

        Returns:
            A new lens that composes this lens with the provided lens/view
        """
        ...

class LensField[B](abc.ABC):
    """Type annotation for lens-enabled fields in dataclasses.

    LensField provides a type-safe way to enable lens access on dataclass fields.
    When a dataclass inherits from HasLensFields, fields annotated with LensField[T]
    can be accessed both as regular attributes and as lenses.

    Type Parameters:
        B: The type of the field value

    Behavior:
        - **Class access** (e.g., `MyClass.field`): Returns a `Lens[MyClass, B]`
          that can be used for functional operations like `get()`, `set()`, `focus()`
        - **Instance access** (e.g., `obj.field`): Returns the actual field value
          of type `B`, behaving like a normal attribute

    Usage:
        For dataclasses deriving from HasLensFields, annotate fields with LensField[T]
        to enable lens access. Regular dataclasses without HasLensFields inheritance
        do not support lens access through field annotations.

    Examples:
        >>> from kups.core.utils.jax import dataclass
        >>> from kups.core.lens import LensField, HasLensFields
        >>> import jax.numpy as jnp
        >>> from jax import Array
        >>>
        >>> @dataclass
        >>> class Point(HasLensFields):
        ...     x: LensField[float]
        ...     y: LensField[Array]
        >>>
        >>> # Instance access - normal field behavior
        >>> point = Point(x=1.0, y=jnp.array([1.0, 2.0, 3.0]))
        >>> point.x  # Returns 1.0
        >>> point.y  # Returns Array([1., 2., 3.])
        >>>
        >>> # Class access - returns a lens
        >>> x_lens = Point.x  # Returns Lens[Point, float]
        >>> x_lens.get(point)  # Returns 1.0
        >>> new_point = x_lens.set(point, 5.0)  # Returns Point(x=5.0, y=...)
        >>>
        >>> # Compose with other lenses
        >>> doubled_x_lens = Point.x.focus(lambda x: x * 2)
        >>> doubled_x_lens.get(point)  # Returns 2.0
        >>>
        >>> # Works in JAX transformations
        >>> @jax.jit
        >>> def increment_y(p: Point) -> Point:
        ...     return Point.y.set(p, p.y + 1.0)

    Notes:
        - Only works with dataclasses that inherit from HasLensFields
        - The metaclass FieldMetaAccess intercepts class attribute access to return lenses
        - Compatible with JAX transformations when used with jax-compatible
          dataclasses (e.g., from `kups.core.utils.jax`)
        - Use `lens_field()` instead of `field()` for type-safe field definitions
          with default values or field options
    """

    def __new__(cls):
        if cls is LensField:
            raise TypeError(
                "Can't instantiate abstract class {}".format(LensField.__name__)
            )
        return super(LensField, cls).__new__(cls)

    @overload
    def __get__[A: HasLensFields](
        self, instance: None, owner: type[A]
    ) -> Lens[A, B]: ...

    @overload
    def __get__(self, instance: Any, owner: type[HasLensFields]) -> B: ...

    @abc.abstractmethod
    def __get__(self, instance, owner) -> Any: ...

    @abc.abstractmethod
    def __set__(self, instance: HasLensFields, value: B): ...

    @abc.abstractmethod
    def __set_name__(self, owner: HasLensFields, name: str) -> None: ...