pickers.dataverse

DataverseSuggestion

Bases: BaseModel

Internal representation of a suggestion entry in Dataverse.

validate_trainer_state classmethod

validate_trainer_state(value)

Validate and convert the trainer_state field from a JSON string to a TrainerState object.

Source code in src/clabe/pickers/dataverse.py

@field_validator("trainer_state", mode="before")
@classmethod
def validate_trainer_state(cls, value):
    """
    Validate and convert the trainer_state field from a JSON string to a TrainerState object.
    """
    if value is None:
        return value
    if isinstance(value, str):
        return TrainerState.model_validate_json(value)
    return value

from_request_output classmethod

from_request_output(
    subject: str, request_output: dict
) -> DataverseSuggestion

Create a _Suggestion instance from a dictionary of data.

Source code in src/clabe/pickers/dataverse.py

@classmethod
def from_request_output(cls, subject: str, request_output: dict) -> "DataverseSuggestion":
    """
    Create a _Suggestion instance from a dictionary of data.
    """
    trainer_state = request_output.get("aibs_trainer_state", None)
    trainer_state = TrainerState.model_validate_json(cls._strip_html(trainer_state)) if trainer_state else None
    return cls(
        subject_id=subject,
        trainer_state=trainer_state,
        task_name=request_output.get("aibs_task_name", None),
        stage_name=request_output.get("aibs_stage_name", None),
        modified_on=request_output.get("modifiedon", None),
        created_on=request_output.get("createdon", None),
    )

from_trainer_state classmethod

from_trainer_state(
    subject: str, trainer_state: TrainerState
) -> DataverseSuggestion

Create a _Suggestion instance from a TrainerState object.

Source code in src/clabe/pickers/dataverse.py

@classmethod
def from_trainer_state(cls, subject: str, trainer_state: TrainerState) -> "DataverseSuggestion":
    """
    Create a _Suggestion instance from a TrainerState object.
    """
    if trainer_state is None:
        raise ValueError("trainer_state cannot be None")
    if trainer_state.stage is None:
        raise ValueError("trainer_state.stage cannot be None")
    return cls(
        subject_id=subject,
        trainer_state=trainer_state,
        task_name=trainer_state.stage.task.name,
        stage_name=trainer_state.stage.name,
    )

DataversePicker

DataversePicker(
    *,
    dataverse_client: Optional[_DataverseRestClient] = None,
    settings: DefaultBehaviorPickerSettings,
    launcher: Launcher,
    ui_helper: Optional[IUiHelper] = None,
    experimenter_validator: Optional[
        Callable[[str], bool]
    ] = validate_username,
    rig_validator: Optional[
        Callable[[Rig], Rig]
    ] = validate_rig_computer_name,
)

Bases: DefaultBehaviorPicker

Picker that integrates with Dataverse to fetch and push trainer state suggestions.

Initializes the DataversePicker.

Parameters:

Name	Type	Description	Default
dataverse_client	Optional[_DataverseRestClient]	Optional Dataverse REST client for making API calls. If not provided, a new client will be created using settings from KeePass.	None
settings	DefaultBehaviorPickerSettings	Settings containing configuration including config_library_dir	required
ui_helper	Optional[IUiHelper]	Helper for user interface interactions	None
experimenter_validator	Optional[Callable[[str], bool]]	Function to validate the experimenter's username. If None, no validation is performed	validate_username
rig_validator	Optional[Callable[[Rig], Rig]]	Function to validate the rig configuration. If None, no validation is performed	validate_rig_computer_name

Source code in src/clabe/pickers/dataverse.py

def __init__(
    self,
    *,
    dataverse_client: Optional[_DataverseRestClient] = None,
    settings: DefaultBehaviorPickerSettings,
    launcher: Launcher,
    ui_helper: Optional[ui.IUiHelper] = None,
    experimenter_validator: Optional[Callable[[str], bool]] = validate_username,
    rig_validator: Optional[Callable[[Rig], Rig]] = validate_rig_computer_name,
):
    """
    Initializes the DataversePicker.

    Args:
        dataverse_client: Optional Dataverse REST client for making API calls. If not provided, a new client will be created using settings from KeePass.
        settings: Settings containing configuration including config_library_dir
        ui_helper: Helper for user interface interactions
        experimenter_validator: Function to validate the experimenter's username. If None, no validation is performed
        rig_validator: Function to validate the rig configuration. If None, no validation is performed
    """
    super().__init__(
        settings=settings,
        launcher=launcher,
        ui_helper=ui_helper,
        experimenter_validator=experimenter_validator,
        rig_validator=rig_validator,
    )
    self._dataverse_client = (
        dataverse_client
        if dataverse_client is not None
        else _DataverseRestClient(_DataverseRestClientSettings.from_keepass())
    )
    self._dataverse_suggestion: Optional[DataverseSuggestion] = None

ui_helper property

ui_helper: IUiHelper

Retrieves the registered UI helper.

Returns:

Name	Type	Description
UiHelper	IUiHelper	The registered UI helper

Raises:

Type	Description
ValueError	If no UI helper is registered

trainer_state property

trainer_state: TrainerState

Returns the current trainer state.

Returns:

Name	Type	Description
TrainerState	TrainerState	The current trainer state

Raises:

Type	Description
ValueError	If the trainer state is not set

session_directory property

session_directory: Path

Returns the directory path for the current session.

session property

session: Session

Returns the current session model.

config_library_dir property

config_library_dir: Path

Returns the path to the configuration library directory.

Returns:

Name	Type	Description
Path	Path	The configuration library directory

rig_dir property

rig_dir: Path

Returns the path to the rig configuration directory.

Returns:

Name	Type	Description
Path	Path	The rig configuration directory

subject_dir property

subject_dir: Path

Returns the path to the subject configuration directory.

Returns:

Name	Type	Description
Path	Path	The subject configuration directory

task_dir property

task_dir: Path

Returns the path to the task configuration directory.

Returns:

Name	Type	Description
Path	Path	The task configuration directory

pick_trainer_state

pick_trainer_state(
    task_model: Type[TTask],
) -> tuple[TrainerState, TTask]

Prompts the user to select or create a trainer state configuration.

Attempts to load trainer state in the following order: 1. If task already exists in launcher, will return an empty TrainerState 2. From subject-specific folder

It will launcher.set_task if the deserialized TrainerState is valid.

Returns:

Name	Type	Description
TrainerState	tuple[TrainerState, TTask]	The deserialized TrainerState object.

Raises:

Type	Description
ValueError	If no valid task file is found.

Source code in src/clabe/pickers/dataverse.py

def pick_trainer_state(self, task_model: Type[TTask]) -> tuple[TrainerState, TTask]:
    """
    Prompts the user to select or create a trainer state configuration.

    Attempts to load trainer state in the following order:
    1. If task already exists in launcher, will return an empty TrainerState
    2. From subject-specific folder

    It will launcher.set_task if the deserialized TrainerState is valid.

    Returns:
        TrainerState: The deserialized TrainerState object.

    Raises:
        ValueError: If no valid task file is found.
    """
    if self._session is None:
        raise ValueError("No session set. Run pick_session first.")
    task_name = task_model.model_fields["name"].default
    if not task_name:
        raise ValueError("Task model does not have a default name.")
    try:
        logger.debug("Attempting to load trainer state dataverse")
        last_suggestions = _get_last_suggestions(self._dataverse_client, self._session.subject, task_name, 1)
    except requests.exceptions.HTTPError as e:
        logger.error("Failed to fetch suggestions from Dataverse: %s", e)
        raise
    except pydantic.ValidationError as e:
        logger.error("Failed to validate suggestion from Dataverse: %s", e)
        raise
    if len(last_suggestions) == 0:
        raise ValueError(
            f"No valid suggestions found in Dataverse for subject {self._session.subject} with task {task_name}."
        )

    _dataverse_suggestion = last_suggestions[0]

    assert _dataverse_suggestion is not None
    if _dataverse_suggestion.trainer_state is None:
        raise ValueError("No trainer state found in the latest suggestion.")
    if _dataverse_suggestion.trainer_state.stage is None:
        raise ValueError("No stage found in the latest suggestion's trainer state.")
    self._dataverse_suggestion = _dataverse_suggestion
    self._trainer_state = _dataverse_suggestion.trainer_state

    assert self._trainer_state is not None
    if not self._trainer_state.is_on_curriculum:
        logging.warning("Deserialized TrainerState is NOT on curriculum.")
    return (
        self.trainer_state,
        task_model.model_validate_json(self.trainer_state.stage.task.model_dump_json()),
    )

push_new_suggestion

push_new_suggestion(trainer_state: TrainerState) -> None

Pushes a new suggestion to Dataverse for the current subject in the launcher. Args: launcher: The Launcher instance containing the current session and subject information. trainer_state: The TrainerState object to be pushed as a new suggestion.

Source code in src/clabe/pickers/dataverse.py

def push_new_suggestion(self, trainer_state: TrainerState) -> None:
    """
    Pushes a new suggestion to Dataverse for the current subject in the launcher.
    Args:
        launcher: The Launcher instance containing the current session and subject information.
        trainer_state: The TrainerState object to be pushed as a new suggestion.
    """
    if self._session is None:
        raise ValueError("No session or subject set in launcher.")

    logger.info("Pushing new suggestion to Dataverse for subject %s", self._session.subject)
    _append_suggestion(self._dataverse_client, self._session.subject, trainer_state)

pick_rig

pick_rig(model: Type[TRig]) -> TRig

Prompts the user to select a rig configuration file.

Searches for available rig configuration files and either automatically selects a single file or prompts the user to choose from multiple options.

Parameters:

Name	Type	Description	Default
model	Type[TRig]	The rig model type to validate against	required

Returns:

Name	Type	Description
TRig	TRig	The selected rig configuration

Raises:

Type	Description
ValueError	If no rig configuration files are found or an invalid choice is made

Source code in src/clabe/pickers/default_behavior.py

def pick_rig(self, model: Type[TRig]) -> TRig:
    """
    Prompts the user to select a rig configuration file.

    Searches for available rig configuration files and either automatically
    selects a single file or prompts the user to choose from multiple options.

    Args:
        model: The rig model type to validate against

    Returns:
        TRig: The selected rig configuration

    Raises:
        ValueError: If no rig configuration files are found or an invalid choice is made
    """
    rig: TRig | None = None
    rig_path: str | None = None

    # Check cache for previously used rigs
    if self._use_cache:
        cache = self._cache_manager.try_get_cache(model.__name__)
    else:
        cache = None

    if cache:
        cache.sort()
        rig_path = self.ui_helper.prompt_pick_from_list(
            cache,
            prompt=f"Choose a rig for {model.__name__}:",
            allow_0_as_none=True,
            zero_label="Select from library",
        )
        if rig_path is not None:
            rig = self._load_rig_from_path(Path(rig_path), model)

    # Prompt user to select a rig if not already selected
    while rig_path is None:
        available_rigs = glob.glob(os.path.join(self.rig_dir, "*.json"))
        # We raise if no rigs are found to prevent an infinite loop
        if len(available_rigs) == 0:
            logger.error("No rig config files found.")
            raise ValueError("No rig config files found.")
        # Use the single available rig config file
        elif len(available_rigs) == 1:
            logger.info("Found a single rig config file. Using %s.", {available_rigs[0]})
            rig_path = available_rigs[0]
            rig = model_from_json_file(rig_path, model)
        else:
            rig_path = self.ui_helper.prompt_pick_from_list(
                available_rigs, prompt=f"Choose a rig for {model.__name__}:"
            )
            if rig_path is not None:
                rig = self._load_rig_from_path(Path(rig_path), model)
    assert rig_path is not None
    assert rig is not None
    if self._rig_validator:
        rig = self._rig_validator(rig)
    # Add the selected rig path to the cache
    self._cache_manager.add_to_cache("rigs", rig_path)
    return rig

pick_session

pick_session(model: Type[TSession] = Session) -> TSession

Prompts the user to select or create a session configuration.

Collects experimenter information, subject selection, and session notes to create a new session configuration with appropriate metadata.

Parameters:

Name	Type	Description	Default
model	Type[TSession]	The session model type to instantiate. Defaults to Session	Session

Returns:

Name	Type	Description
TSession	TSession	The created or selected session configuration

Source code in src/clabe/pickers/default_behavior.py

def pick_session(self, model: Type[TSession] = Session) -> TSession:
    """
    Prompts the user to select or create a session configuration.

    Collects experimenter information, subject selection, and session notes
    to create a new session configuration with appropriate metadata.

    Args:
        model: The session model type to instantiate. Defaults to Session

    Returns:
        TSession: The created or selected session configuration
    """

    experimenter = self.prompt_experimenter(strict=True)
    subject = self.choose_subject(self.subject_dir)

    if not (self.subject_dir / subject).exists():
        logger.info("Directory for subject %s does not exist. Creating a new one.", subject)
        os.makedirs(self.subject_dir / subject)

    notes = self.ui_helper.prompt_text("Enter notes: ")
    session = model(
        subject=subject,
        notes=notes,
        experimenter=experimenter if experimenter is not None else [],
        commit_hash=self._launcher.repository.head.commit.hexsha,
        allow_dirty_repo=self._launcher.settings.debug_mode or self._launcher.settings.allow_dirty,
        skip_hardware_validation=self._launcher.settings.skip_hardware_validation,
    )
    self._session = session
    return session

pick_task

pick_task(model: Type[TTask]) -> TTask

Prompts the user to select or create a task configuration.

Attempts to load task in the following order: 1. From CLI if already set 2. From subject-specific folder 3. From user selection in task library

Parameters:

Name	Type	Description	Default
model	Type[TTask]	The task model type to validate against	required

Returns:

Name	Type	Description
TTask	TTask	The created or selected task configuration

Raises:

Type	Description
ValueError	If no valid task file is found

Source code in src/clabe/pickers/default_behavior.py

def pick_task(self, model: Type[TTask]) -> TTask:
    """
    Prompts the user to select or create a task configuration.

    Attempts to load task in the following order:
    1. From CLI if already set
    2. From subject-specific folder
    3. From user selection in task library

    Args:
        model: The task model type to validate against

    Returns:
        TTask: The created or selected task configuration

    Raises:
        ValueError: If no valid task file is found
    """
    task: Optional[TTask] = None
    if self._session is None:
        raise ValueError("Session must be picked (pick_session) before picking task.")

    try:
        f = self.subject_dir / self._session.subject / (ByAnimalFiles.TASK.value + ".json")
        logger.info("Attempting to load task from subject folder: %s", f)
        task = model_from_json_file(f, model)
    except (ValueError, FileNotFoundError, pydantic.ValidationError) as e:
        logger.warning("Failed to find a valid task file. %s", e)
    else:
        logger.info("Found a valid task file in subject folder!")
        _is_manual = not self.ui_helper.prompt_yes_no_question("Would you like to use this task?")
        if not _is_manual:
            if task is not None:
                return task
            else:
                logger.error("No valid task file found in subject folder.")
                raise ValueError("No valid task file found.")
        else:
            task = None

    # If not found, we prompt the user to choose/enter a task file
    while task is None:
        try:
            _path = Path(os.path.join(self.config_library_dir, self.task_dir))
            available_files = glob.glob(os.path.join(_path, "*.json"))
            if len(available_files) == 0:
                break
            path = self.ui_helper.prompt_pick_from_list(
                available_files, prompt=f"Choose a task for {model.__name__}:"
            )
            if not isinstance(path, str):
                raise ValueError("Invalid choice.")
            if not os.path.isfile(path):
                raise FileNotFoundError(f"File not found: {path}")
            task = model_from_json_file(path, model)
            logger.info("User entered: %s.", path)
        except pydantic.ValidationError as e:
            logger.error("Failed to validate pydantic model. Try again. %s", e)
        except (ValueError, FileNotFoundError) as e:
            logger.info("Invalid choice. Try again. %s", e)
    if task is None:
        logger.error("No task file found.")
        raise ValueError("No task file found.")

    return task

choose_subject

choose_subject(directory: str | PathLike) -> str

Prompts the user to select or manually enter a subject name.

Allows the user to either type a new subject name or select from existing subject directories.

Parameters:

Name	Type	Description	Default
directory	str | PathLike	Path to the directory containing subject folders	required

Returns:

Name	Type	Description
str	str	The selected or entered subject name.

Example

# Choose a subject from the subjects directory
subject = picker.choose_subject("Subjects")

Source code in src/clabe/pickers/default_behavior.py

def choose_subject(self, directory: str | os.PathLike) -> str:
    """
    Prompts the user to select or manually enter a subject name.

    Allows the user to either type a new subject name or select from
    existing subject directories.

    Args:
        directory: Path to the directory containing subject folders

    Returns:
        str: The selected or entered subject name.

    Example:
        ```python
        # Choose a subject from the subjects directory
        subject = picker.choose_subject("Subjects")
        ```
    """
    if self._use_cache:
        subjects = self._cache_manager.try_get_cache("subjects")
    else:
        subjects = None
    if subjects:
        subjects.sort()
        subject = self.ui_helper.prompt_pick_from_list(
            subjects,
            prompt="Choose a subject:",
            allow_0_as_none=True,
            zero_label="Enter manually",
        )
    else:
        subject = None

    while subject is None:
        subject = self.ui_helper.input("Enter subject name: ")
        if subject == "":
            subject = None
    self._cache_manager.add_to_cache("subjects", subject)
    return subject

prompt_experimenter

prompt_experimenter(
    strict: bool = True,
) -> Optional[List[str]]

Prompts the user to enter the experimenter's name(s).

Accepts multiple experimenter names separated by commas or spaces. Validates names using the configured validator function if provided.

Parameters:

Name	Type	Description	Default
strict	bool	Whether to enforce non-empty input. Defaults to True	True

Returns:

Type	Description
Optional[List[str]]	Optional[List[str]]: List of experimenter names

Example

# Prompt for experimenter with validation
names = picker.prompt_experimenter(strict=True)
print("Experimenters:", names)

Source code in src/clabe/pickers/default_behavior.py

def prompt_experimenter(self, strict: bool = True) -> Optional[List[str]]:
    """
    Prompts the user to enter the experimenter's name(s).

    Accepts multiple experimenter names separated by commas or spaces.
    Validates names using the configured validator function if provided.

    Args:
        strict: Whether to enforce non-empty input. Defaults to True

    Returns:
        Optional[List[str]]: List of experimenter names

    Example:
        ```python
        # Prompt for experimenter with validation
        names = picker.prompt_experimenter(strict=True)
        print("Experimenters:", names)
        ```
    """
    if self._use_cache:
        experimenters_cache = self._cache_manager.try_get_cache("experimenters")
    else:
        experimenters_cache = None
    experimenter: Optional[List[str]] = None
    _picked: str | None = None
    while experimenter is None:
        if experimenters_cache:
            experimenters_cache.sort()
            _picked = self.ui_helper.prompt_pick_from_list(
                experimenters_cache,
                prompt="Choose an experimenter:",
                allow_0_as_none=True,
                zero_label="Enter manually",
            )
        if _picked is None:
            _input = self.ui_helper.prompt_text("Experimenter name: ")
        else:
            _input = _picked
        experimenter = _input.replace(",", " ").split()
        if strict & (len(experimenter) == 0):
            logger.info("Experimenter name is not valid. Try again.")
            experimenter = None
        else:
            if self._experimenter_validator:
                for name in experimenter:
                    if not self._experimenter_validator(name):
                        logger.warning("Experimenter name: %s, is not valid. Try again", name)
                        experimenter = None
                        break
    self._cache_manager.add_to_cache("experimenters", ",".join(experimenter))
    return experimenter

dump_model

dump_model(
    model: Union[Rig, Task, TrainerState],
) -> Optional[Path]

Saves the provided model to the appropriate configuration file.

Parameters:

Name	Type	Description	Default
model	Union[Rig, Task, TrainerState]	The model instance to save	required

Returns:

Type	Description
Optional[Path]	Optional[Path]: The path to the saved model file, or None if not saved

Source code in src/clabe/pickers/default_behavior.py

def dump_model(
    self,
    model: Union[Rig, Task, TrainerState],
) -> Optional[Path]:
    """
    Saves the provided model to the appropriate configuration file.

    Args:
        model: The model instance to save

    Returns:
        Optional[Path]: The path to the saved model file, or None if not saved
    """

    path: Path
    if isinstance(model, Rig):
        path = self.rig_dir / ("rig.json")
    elif isinstance(model, Task):
        if self._session is None:
            raise ValueError("Session must be picked (pick_session) before dumping task.")
        path = Path(self.subject_dir) / self._session.subject / (ByAnimalFiles.TASK.value + ".json")
    elif isinstance(model, TrainerState):
        if self._session is None:
            raise ValueError("Session must be picked (pick_session) before dumping trainer state.")
        path = Path(self.subject_dir) / self._session.subject / (ByAnimalFiles.TRAINER_STATE.value + ".json")
    else:
        raise ValueError("Model type not supported for dumping.")

    os.makedirs(path.parent, exist_ok=True)
    if path.exists():
        overwrite = self.ui_helper.prompt_yes_no_question(f"File {path} already exists. Overwrite?")
        if not overwrite:
            logger.info("User chose not to overwrite the existing file: %s", path)
            return None
    with open(path, "w", encoding="utf-8") as f:
        f.write(model.model_dump_json(indent=2))
        logger.info("Saved model to %s", path)
    return path