Python Decorator to convert Pandas DataFrame to dict

How can I create a Python decorator that can wrap any function that returns a Pandas Dataframe, and convert this returned Dataframe into a dictionary of the format:

output = {col: list(df[col]) for col in df.columns}

For example:

import pandas as pd

def some_func_that_returns_df():
  df = pd.DataFrame([1,2,3,4,5], columns=['a'])
  return df


def wrapper(df: pd.DataFrame):
    _df = {col: df[col].to_list() for col in df.columns}
    return _df

df = wrapper(some_func_that_returns_df()) 
print(df)
>>> {'a': [1, 2, 3, 4, 5]}

However, I would like to achieve this:

@df_to_dict
def some_func_that_returns_df():
  df = pd.DataFrame([1,2,3,4,5], columns=['a'])
  return df

df_as_dict = some_func_that_returns_df()

Answer

You can use the following:

def df_to_dict(func):
    def to_dict():
        return func().to_dict(orient='list')
    return to_dict

@df_to_dict
def some_func_that_returns_df():
    df = pd.DataFrame([1,2,3,4,5], columns=['a'])
    return df

output:

>>> some_func_that_returns_df()
{'a': [1, 2, 3, 4, 5]}