Skip to content

pickle_writer

Pickle writing functionality for the devices_rap pipeline.

create_pickle(output_workbooks, output_directory, fin_month, fin_year)

Create a pickle file containing the processed data for all regions. The pickle file will be saved in the output directory with a name that includes the financial year and month.

Parameters:

Name Type Description Default
output_workbooks dict

The processed data for each region

required
output_directory Path

The path to save the pickle file to

required
fin_month str

The financial month for which the data is being processed

required
fin_year str

The financial year for which the data is being processed

required

Returns:

Type Description
None
Source code in devices_rap/data_io/output/pickle_writer.py
def create_pickle(
    output_workbooks: dict[str, dict[str, pd.DataFrame]],
    output_directory: Path,
    fin_month: str,
    fin_year: str,
):
    """
    Create a pickle file containing the processed data for all regions. The pickle file will be saved
    in the output directory with a name that includes the financial year and month.

    Parameters
    ----------
    output_workbooks : dict
        The processed data for each region
    output_directory : Path
        The path to save the pickle file to
    fin_month : str
        The financial month for which the data is being processed
    fin_year : str
        The financial year for which the data is being processed

    Returns
    -------
    None
    """
    logger.info("Creating pickle file")

    output_file = output_directory / f"{fin_year}_{fin_month}_amber_report_all_regions.pkl"
    if output_file.exists():
        logger.warning(f"Overwriting the existing pickle file: {output_file}")
        output_file.unlink()
    else:
        logger.info(f"Creating pickle file for all regions for {fin_month} {fin_year}")

    with open(output_file, "wb") as f:
        pd.to_pickle(output_workbooks, f)

    logger.success(f"Pickle file for all regions for {fin_month} {fin_year} created successfully.")