Use Custom Polars Expression For Returning Dataframes, Where Applicable

[AFg6K7h4fhy2]

In this comment, #33 (comment), DHM references the following code (which is from recode_locations.py):

def loc_flusight_code_to_abbr(
    df: pl.DataFrame, location_col: str
) -> pl.DataFrame:
    """
    Takes the location columns of a Polars
    dataframe (formatted as FluSight codes for
    US two-letter jurisdictions) and recodes
    it to US jurisdictional abbreviations,
    using location_table, which is a Polars
    dataframe contained in forecasttools.

    Parameters
    ----------
    df
        A Polars dataframe with a location
        column consisting of US
        jurisdictional FluSight codes.
    location_col
        The name of the dataframe's location
        column.

    Returns
    -------
    pl.DataFrame
        A Polars dataframe with the location
        column formatted as US two-letter
        jurisdictional abbreviations.
    """
    # get forecasttools location table
    loc_table = forecasttools.location_table
    # recode location codes to location abbreviations
    loc_recoded_df = df.with_columns(
        pl.col(location_col).replace(
            old=loc_table["location_code"], new=loc_table["short_name"]
        )
    )
    return loc_recoded_df

DHM writes:

For now it is fine to have these as functions that take and return dataframes, more polars idiomatic to define a custom expression that accomplishes this https://docs.pola.rs/api/python/stable/reference/api.html.

This issue covers adding this functionality to forecasttools. The author adds this issue to be the next Sprint since it seems doable in that time frame.