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

import argschema
import pathlib
from marshmallow import post_load

from allensdk.brain_observatory.vbn_2022.utils.schemas import (
    ProbeToSkip)


class VBN2022InputJsonWriterSchema(argschema.ArgSchema):

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

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

    json_output_dir = argschema.fields.OutputDir(
            required=True,
            description=("Directory where input JSONs will be written"))

    nwb_output_dir = argschema.fields.OutputDir(
            required=True,
            description=("Directory where NWB files will be written"))

    clobber = argschema.fields.Boolean(
            default=False,
            description=("If false, throw an error if output files "
                         "already exist"))

    json_prefix = argschema.fields.Str(
            required=False,
            default='vbn_ecephys_session',
            allow_none=False,
            description=('The files written by this module will be '
                         'named like '
                         '{json_prefix}_{session_id}_input.json'))

    nwb_prefix = argschema.fields.Str(
            required=False,
            default='ecepys',
            allow_none=False,
            description=('The NWB files specified in the input JSONs '
                         'will be named like '
                         '{nwb_prefix}_{session_id}.nwb'))

    @post_load
    def create_path_lookup(self, data, **kwargs):
        """
        Construct lookups mapping ecephys_session_id to
        the input_json_path and the nwb_file_path
        """
        json_lookup = dict()
        nwb_lookup = dict()

        unq_json = set()
        unq_nwb = set()

        json_dir_path = pathlib.Path(data['json_output_dir'])
        nwb_dir_path = pathlib.Path(data['nwb_output_dir'])

        for ecephys_id in data['ecephys_session_id_list']:
            json_name = f"{data['json_prefix']}_{ecephys_id}_input.json"
            if json_name in unq_json:
                raise RuntimeError("This configuration would write "
                                   f"{json_name} more than once")
            unq_json.add(json_name)
            json_path = json_dir_path / json_name
            if not data['clobber'] and json_path.is_file():
                raise RuntimeError(f"{json_path.resolve().absolute()} "
                                   "already exists; "
                                   "run with clobber=True to overwrite")
            json_lookup[ecephys_id] = json_path

            nwb_name = f"{data['nwb_prefix']}_{ecephys_id}.nwb"
            if nwb_name in unq_nwb:
                raise RuntimeError("This configuration would write "
                                   f"{nwb_name} more than once")

            unq_nwb.add(nwb_name)
            nwb_path = nwb_dir_path / f'{ecephys_id}' / nwb_name
            nwb_lookup[ecephys_id] = nwb_path

        data['json_path_lookup'] = json_lookup
        data['nwb_path_lookup'] = nwb_lookup

        return data