Source code for ursa.temporal

"""Ursa subclasses of the core ``temporaldata`` primitives.

These classes shadow the upstream names so callers do
``from ursa import Data, IrregularTimeSeries, RegularTimeSeries`` and get
drop-in replacements that carry catalog-row metadata. With ``metadata=None``
each class is behaviorally identical to its upstream parent.

Future cloud-streaming subclasses (R2 Zarr / Lance — ENG-899) will subclass
*these* classes, inheriting the metadata slot for free.

Storage convention follows upstream's pattern (``_sampling_rate`` →
``sampling_rate`` property): the row is stored as the private attribute
``_metadata`` so it bypasses ``ArrayDict.__setattr__``'s ndarray check, and
exposed via a read-only ``metadata`` property. The single ``slice``
override per class is needed because upstream's ``slice`` builds the
result via ``__class__.__new__`` and only copies the private attrs it
knows about (``_sampling_rate``, ``_timekeys``, ``_domain``).
"""

from __future__ import annotations

from typing import Any, cast

import numpy as np
from temporaldata import Data as _Data
from temporaldata import IrregularTimeSeries as _IrregularTimeSeries
from temporaldata import RegularTimeSeries as _RegularTimeSeries

from ursa.catalog.schemas import ModalityRow, RecordingRow




def _check_metadata(value: Any, expected_type: type, owner_name: str) -> None:
    if value is not None and not isinstance(value, expected_type):
        raise TypeError(
            f"{owner_name}.metadata must be a {expected_type.__name__} or None, "
            f"got {type(value).__name__}"
        )





class Data(_Data):  # type: ignore[misc]
    """``temporaldata.Data`` carrying optional :class:`RecordingRow` metadata."""

    def __init__(
        self,
        *,
        metadata: RecordingRow | None = None,
        **kwargs: Any,
    ) -> None:
        _check_metadata(metadata, RecordingRow, "Data")
        super().__init__(**kwargs)
        self._metadata = metadata

    @property
    def metadata(self) -> RecordingRow | None:
        return self._metadata



    def slice(self, start: float, end: float, reset_origin: bool = True) -> "Data":
        out = cast("Data", super().slice(start, end, reset_origin))
        out._metadata = self._metadata
        return out






class IrregularTimeSeries(_IrregularTimeSeries):  # type: ignore[misc]
    """``temporaldata.IrregularTimeSeries`` carrying optional :class:`ModalityRow` metadata."""

    def __init__(
        self,
        timestamps: np.ndarray,
        *,
        metadata: ModalityRow | None = None,
        **kwargs: Any,
    ) -> None:
        _check_metadata(metadata, ModalityRow, "IrregularTimeSeries")
        super().__init__(timestamps, **kwargs)
        self._metadata = metadata

    @property
    def metadata(self) -> ModalityRow | None:
        return self._metadata



    def slice(self, start: float, end: float, reset_origin: bool = True) -> "IrregularTimeSeries":
        out = cast("IrregularTimeSeries", super().slice(start, end, reset_origin))
        out._metadata = self._metadata
        return out






class RegularTimeSeries(_RegularTimeSeries):  # type: ignore[misc]
    """``temporaldata.RegularTimeSeries`` carrying optional :class:`ModalityRow` metadata."""

    def __init__(
        self,
        *,
        metadata: ModalityRow | None = None,
        **kwargs: Any,
    ) -> None:
        _check_metadata(metadata, ModalityRow, "RegularTimeSeries")
        super().__init__(**kwargs)
        self._metadata = metadata

    @property
    def metadata(self) -> ModalityRow | None:
        return self._metadata



    def slice(self, start: float, end: float, reset_origin: bool = True) -> "RegularTimeSeries":
        out = cast("RegularTimeSeries", super().slice(start, end, reset_origin))
        out._metadata = self._metadata
        return out




__all__ = ["Data", "IrregularTimeSeries", "RegularTimeSeries"]