modules/utils.py

from pandas import DataFrame


def _get_nested_metadata(metadata_key: str, metadata: dict[dict], metadata_value: str):
    """Get nested metadata from `dict[dict, Any]` structure.
    Returns "metadata_key" as default value.
    """
    getter = metadata.get(metadata_key, {})
    return getter.get(metadata_value, metadata_key)


def get_agency_metadata_values(
        df: DataFrame, 
        agency_column: str, 
        metadata: dict,         
        metadata_value: str, 
    ):
    """Get a specific value from agency metadata (e.g., get acronym for Department of Homeland Security).

    Args:
        df (DataFrame): Input data.
        agency_column (str): Column containing agency identifier.
        metadata (dict): Agency metadata.
        metadata_value (str): Value of interest from agency metadata.

    Returns:
        pd.Series: Pandas Series of new values for adding to DataFrame.
    """    
    if metadata_value == "acronym":
        metadata_value = "short_name"
    return df.loc[:, agency_column].apply(
        lambda x: [
            _get_nested_metadata(x_item, metadata=metadata, metadata_value=metadata_value)
            for x_item 
            in x
            ] if isinstance(x, list) else _get_nested_metadata(x, metadata=metadata, metadata_value=metadata_value)
        )