Source code for allensdk.brain_observatory.vbn_2022.metadata_writer.schemas

import pathlib

import argschema
from allensdk.brain_observatory.vbn_2022.utils.schemas import ProbeToSkip
from allensdk.brain_observatory.behavior.behavior_project_cache.project_metadata_writer.schemas import BaseMetadataWriterInputSchema  # noqa: E501
from marshmallow import post_load




class VBN2022MetadataWriterInputSchema(BaseMetadataWriterInputSchema):

    ecephys_session_id_list = argschema.fields.List(
        argschema.fields.Int,
        required=True,
        description=(
            "List of ecephys_sessions.id values " "of sessions to be released"
        ),
    )

    failed_ecephys_session_id_list = argschema.fields.List(
        argschema.fields.Int,
        required=False,
        default=None,
        allow_none=True,
        description=(
            "List of ecephys_sessions.id values "
            "associated with this release that were "
            "failed. These are required to "
            "self-consistently construct the history of "
            "each mouse passing through the apparatus."
        ),
    )

    probes_to_skip = argschema.fields.List(
        argschema.fields.Nested(ProbeToSkip),
        required=False,
        default=None,
        allow_none=True,
        description=("List of probes to skip"),
    )

    ecephys_nwb_dir = argschema.fields.InputDir(
        required=True,
        allow_none=False,
        description=(
            "The directory where ecephys_nwb sessions are " "to be found"
        ),
    )

    ecephys_nwb_prefix = argschema.fields.Str(
        required=False,
        default="ecephys_session",
        description=(
            "Ecephys session NWB files will be looked for "
            "in the form "
            "{ecephys_nwb_dir}/{ecephys_nwb_prefix}_{ecephys_session_id}.nwb"
        ),
    )

    supplemental_data = argschema.fields.List(
        argschema.fields.Dict,
        default=None,
        allow_none=True,
        description=(
            "List of dicts definining any supplemental columns "
            "that need to be added to the ecephys_sessions.csv "
            "table. Each dict should represent a row in a dataframe "
            "that will get merged on ecephys_session_id with "
            "the ecephys_sessions table (row must therefore contain "
            "ecephys_session_id)"
        ),
    )
    n_workers = argschema.fields.Int(
        default=8,
        allow_none=True,
        description='Number of workers for reading from pkl file. '
                    'Default=8 due to issues with making too many '
                    'requests to the database. Increase if too slow, decrease '
                    'if the database rejects the connection'
    )



    @post_load
    def validate_paths(self, data, **kwargs):
        fname_lookup = {
            "units_path": "units.csv",
            "channels_path": "channels.csv",
            "probes_path": "probes.csv",
            "ecephys_sessions_path": "ecephys_sessions.csv",
            "behavior_sessions_path": "behavior_sessions.csv",
        }

        out_dir = pathlib.Path(data["output_dir"])
        msg = ""
        for fname_k in fname_lookup.keys():
            full_path = out_dir / fname_lookup[fname_k]
            if full_path.exists() and not data["clobber"]:
                msg += f"{full_path.resolve().absolute()}\n"
            data[fname_k] = str(full_path.resolve().absolute())

        if len(msg) > 0:
            raise RuntimeError(
                "The following files already exist\n"
                f"{msg}"
                "Run with clobber=True if you want to overwrite"
            )
        return data