stress_at_work_analysis/features/light.py

from collections.abc import Collection

import pandas as pd

from config.models import LightSensor, Participant
from setup import db_engine, session


def get_light_data(usernames: Collection) -> pd.DataFrame:
    """
    Read the data from the light sensor table and return it in a dataframe.

    Parameters
    ----------
    usernames: Collection
        A list of usernames to put into the WHERE condition.

    Returns
    -------
    df_light: pd.DataFrame
        A dataframe of light data.
    """
    query_light = (
        session.query(LightSensor, Participant.username)
        .filter(Participant.id == LightSensor.participant_id)
        .filter(Participant.username.in_(usernames))
    )
    with db_engine.connect() as connection:
        df_light = pd.read_sql(query_light.statement, connection)
    return df_light