Python 装饰器将 Pandas DataFrame 转换为 dict

Question

我如何创建一个 Python 装饰器，它可以包装 returns 一个 Pandas Dataframe 的任何函数，并将这个返回的 Dataframe 转换成格式的字典：

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

例如：

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]}

不过，我想实现这个：

@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 1

您的装饰器必须 return 一个函数。

def df_to_dict_decorator(func_that_returns_df):
    def func_that_returns_dict(*args, **kwargs):
        return wrapper(func_that_returns_df(*args, **kwargs)) # this uses your "wrapper" function, consider renaming it to df_to_dict
    return func_that_returns_dict

@df_to_dict_decorator
def foo(df):
    return df

演示：

>>> foo(pd.DataFrame([1, 2, 3]))
{0: [1, 2, 3]}

另请注意，DataFrames 已经有一个 .to_dict 方法。

Answer 2

您可以使用以下内容：

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

输出：

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

Python 装饰器将 Pandas DataFrame 转换为 dict

Python Decorator to convert Pandas DataFrame to dict

python

pandas

python-decorators