src/utils/s3/decorators.py

import pandas as pd

from .save_file_to_s3 import save_csv_to_s3


def save_output_csv_to_s3(save_path: str, bucket: str = "lebesgue-common-bucket"):
    def outer_wrapper(func):
        def wrapper(*args, **kwargs):

            output = func(*args, **kwargs)

            if type(output) == pd.DataFrame:
                save_csv_to_s3(df=output, path=save_path, index=False, bucket=bucket)
            else:
                [
                    save_csv_to_s3(df=df, path=path, index=False, bucket=bucket)
                    for df, path in zip(output, save_path)
                ]

            return output

        return wrapper

    return outer_wrapper