session_pipeline¶

def create_processors(
    dataset: Dataset,
    *,
    strict_parsing: bool = False,
    include: Sequence[str] = (),
    exclude: Sequence[str] = (),
) -> list[AbstractProcessor]:
    """Return the ordered processor list for *dataset*, dispatching on version.

    Parameters
    ----------
    dataset:
        The loaded contraqctor Dataset. Its ``.version`` attribute determines
        which processor variants are selected.
    strict_parsing:
        Passed through to every processor. When ``True``, any parsing anomaly
        raises; when ``False`` (default) it logs a warning and continues.
    include:
        If non-empty, keep only processors whose ``output_name`` is listed.
    exclude:
        Drop processors whose ``output_name`` is listed. Applied after
        *include*, so an name in both is dropped.

    Returns
    -------
    list[AbstractProcessor]
        Processors in the order they must be applied; ``session`` is always
        first and is never filtered out — it carries the session's identity, so
        every other table would lose its join key without it.
        :class:`~.processing.SessionMetadataProcessor` is also unconditional in
        the sense that it needs no arguments: it derives the session root from
        the dataset's own Session stream.
    """

    processors: list[AbstractProcessor] = [
        SessionMetadataProcessor(dataset, strict_parsing=strict_parsing),
        resolve_position_velocity_processor(dataset, strict_parsing=strict_parsing),
        resolve_site_table_processor(dataset, strict_parsing=strict_parsing),
        LicksProcessor(dataset, strict_parsing=strict_parsing),
        SniffingProcessor(dataset, strict_parsing=strict_parsing),
        SoftwareEventsProcessor(dataset, strict_parsing=strict_parsing),
        EventsProcessor(dataset, strict_parsing=strict_parsing),
    ]
    return filter_processors(processors, include=include, exclude=exclude)

def filter_processors(
    processors: Sequence[AbstractProcessor],
    *,
    include: Sequence[str] = (),
    exclude: Sequence[str] = (),
) -> list[AbstractProcessor]:
    """Select processors by ``output_name``, always keeping ``session``.

    Empty *include* means "keep everything"; *exclude* is applied second.
    ``session`` survives both, because dropping it would leave every other
    table without the identity row it joins to.
    """
    include_set, exclude_set = frozenset(include), frozenset(exclude)

    def _keep(proc: AbstractProcessor) -> bool:
        name = proc.output_name
        if name == SessionMetadataProcessor.__output_name__:
            return True
        if include_set and name not in include_set:
            logger.debug("skip %s (not in include list)", name)
            return False
        if name in exclude_set:
            logger.debug("skip %s (excluded)", name)
            return False
        return True

    return [p for p in processors if _keep(p)]

def resolve_site_table_processor(
    dataset: Dataset,
    *,
    strict_parsing: bool = False,
) -> SiteTableProcessor | LegacySiteTableProcessor:
    """Return the correct site-table processor for *dataset*'s version."""
    version = semver.Version.parse(str(dataset.version))
    cls = LegacySiteTableProcessor if version < _LEGACY_VERSION_CUTOFF else SiteTableProcessor
    return cls(dataset, strict_parsing=strict_parsing)

def resolve_position_velocity_processor(
    dataset: Dataset,
    *,
    sampling_rate_hz: float | None = 250.0,
    strict_parsing: bool = False,
) -> PositionAndVelocityProcessor | LegacyPositionAndVelocityProcessor:
    """Return the correct position/velocity processor for *dataset*'s version."""
    version = semver.Version.parse(str(dataset.version))
    cls = LegacyPositionAndVelocityProcessor if version < _LEGACY_VERSION_CUTOFF else PositionAndVelocityProcessor
    return cls(dataset, sampling_rate_hz=sampling_rate_hz, strict_parsing=strict_parsing)

def process_session(
    dataset: Dataset | Path | str,
    output_dir: Path | str = ".",
    *,
    strict_parsing: bool = False,
    include: Sequence[str] = (),
    exclude: Sequence[str] = (),
    processors: Sequence[AbstractProcessor] | None = None,
    on_error: Callable[[AbstractProcessor, Exception], None] | None = None,
    write_parquet: bool = True,
    write_nwb: bool = False,
) -> dict[str, pd.DataFrame]:
    """Run every processor and write the outputs chosen by *write_parquet* / *write_nwb*.

    Each processor's ``output_name`` attribute determines its parquet filename,
    e.g. ``sites.parquet``, ``position_velocity.parquet``, etc.

    The two output formats are independent switches over the same computed
    frames. Every processor runs regardless — the flags choose what reaches disk,
    not what is computed — so the returned dict is the same either way.

    Log lines are prefixed with the session id, taken from the dataset, so
    per-processor progress stays grep-able when many sessions run in one batch.

    Parameters
    ----------
    dataset:
        A loaded contraqctor Dataset, or the path to a raw session directory to
        load one from. Its version determines which processor variants are
        selected (legacy vs current).
    output_dir:
        Directory where outputs are written; defaults to the current working
        directory. Created if absent, unless both *write_parquet* and
        *write_nwb* are ``False``, in which case nothing is written and no
        directory is made.
    strict_parsing:
        Passed to all processors.
    include, exclude:
        Processor ``output_name`` filters, forwarded to
        :func:`create_processors`. ``session`` is never filtered out. Ignored
        when *processors* is given, since that list is already final.
    processors:
        Use this exact, already-constructed processor list instead of calling
        :func:`create_processors` internally — the escape hatch for a custom or
        third-party processor. *strict_parsing*, *include* and *exclude* are
        then irrelevant to processor construction. ``None`` (default) builds
        the list from the dataset.
    on_error:
        Called as ``on_error(processor, exception)`` when a processor's
        ``compute()`` raises, instead of letting the exception propagate.
        The callback decides what happens next: returning normally skips that
        processor and continues with the rest; re-raising aborts the run
        immediately (the callback's own raised exception propagates out of
        this function). ``None`` (default) means an exception propagates
        immediately, same as before this parameter existed — every existing
        caller keeps its current behavior unchanged.
    write_parquet:
        When ``True`` (default), write one ``output_dir/{output_name}.parquet``
        per processor, with provenance promoted into the parquet schema. Set to
        ``False`` to compute without touching disk — the frames still come back
        in the return value.
    write_nwb:
        When ``True``, write ``output_dir/{session_id}.nwb.zarr`` from the same
        processor list, so one filtered selection can produce both output
        formats. Requires the AIND metadata JSON files in the session root; a
        session missing them fails the NWB step, and that failure propagates
        like any other. Defaults to ``False``.

    Returns
    -------
    dict[str, pd.DataFrame]
        DataFrames for every processor that computed successfully, keyed by
        ``output_name`` — independent of which formats were written. A processor
        whose ``compute()`` raised and whose failure was absorbed by *on_error*
        (rather than re-raised) is simply absent from this dict.
    """
    if isinstance(dataset, (str, Path)):
        from aind_behavior_vr_foraging.data_contract import dataset as load_dataset

        dataset = load_dataset(Path(dataset))

    output_dir = Path(output_dir)
    if write_parquet or write_nwb:
        output_dir.mkdir(parents=True, exist_ok=True)

    root = session_root(dataset)
    selected = (
        processors
        if processors is not None
        else create_processors(dataset, strict_parsing=strict_parsing, include=include, exclude=exclude)
    )

    all_data: dict[str, pd.DataFrame] = {}
    for proc in selected:
        name = proc.output_name
        logger.info("[%s] compute: %s → %s", root.name, proc.__class__.__name__, name)
        try:
            # compute() stamps provenance attrs automatically (see AbstractProcessor.compute)
            df = proc.compute()
        except Exception as exc:
            if on_error is None:
                raise
            on_error(proc, exc)
            continue
        all_data[name] = df
        if write_parquet:
            _write_parquet(df, output_dir / f"{name}.parquet")
            logger.info("[%s]   saved %d rows → %s.parquet", root.name, len(df), name)
        else:
            logger.info("[%s]   %d rows (parquet skipped)", root.name, len(df))

    if write_nwb:
        _write_nwb_zarr(dataset, root, output_dir, selected)

    return all_data