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) )