Python如何用装饰器包装pandas函数实现KeyError时自动重置索引

代码库中常见如下pandas处理逻辑，通过链式pipe调用批量给DataFrame关联补充统计字段：

df # 包含userId、sessionId列的pandas DataFrame

def add_session_statistics(df):
   df_statistics = get_session_statistics(df.sessionId.unique())
   return df.merge(df_statistics, on='sessionId', how='left')

# 注：原示例此处存在两处笔误：内部递归调用自身、merge关联键写错，实际应为调用get_user_statistics、关联键为userId
def add_user_statistics(df):
   df_statistics = get_user_statistics(df.userId.unique())
   return df.merge(df_statistics, on='userId', how='left')

# 链式调用做特征扩充
df_enriched = (df
               .pipe(add_session_statistics)
               .pipe(add_user_statistics)
               )

但代码库其他逻辑中，会把userId、sessionId设置为DataFrame的多级索引：

X = df.set_index(['userId', 'sessionId'])

这种场景下如果不手动先执行reset_index()，直接在X上调用add_{something}_statistics()系列函数，会因为关联键不存在触发KeyError。需要实现一个通用装饰器，让这类函数遇到上述KeyError时，自动重置索引完成处理逻辑，兼容两种DataFrame结构。

写一个通用装饰器即可，核心逻辑如下：

• 先尝试直接执行原处理函数
• 若捕获到KeyError，先判断报错的缺失键是否存在于DataFrame的索引层级名中
• 确认是索引导致的键缺失时，自动重置索引将索引层级转为普通列，再执行原处理逻辑
• 非索引原因导致的KeyError会正常抛出，不会吞掉原有报错

装饰器代码如下：

from functools import wraps
import pandas as pd

def auto_reset_index(func):
    @wraps(func)
    def wrapper(df: pd.DataFrame, *args, **kwargs):
        try:
            # 拷贝原对象避免修改输入数据
            return func(df.copy(), *args, **kwargs)
        except KeyError as e:
            missing_keys = e.args[0]
            # 兼容单键、多键缺失的报错场景
            if isinstance(missing_keys, str):
                missing_keys = [missing_keys]
            # 校验所有缺失键是否都属于索引层级
            if all(key in df.index.names for key in missing_keys):
                df_reset = df.reset_index()
                res = func(df_reset, *args, **kwargs)
                # 如果需要处理后保留原索引结构，取消下一行注释即可
                # res = res.set_index(list(df.index.names))
                return res
            # 非索引导致的键缺失正常抛出报错
            raise
    return wrapper

给所有add_{something}_statistics系列函数加上装饰器即可，不需要修改原有函数内部逻辑：

@auto_reset_index
def add_session_statistics(df):
   df_statistics = get_session_statistics(df.sessionId.unique())
   return df.merge(df_statistics, on='sessionId', how='left')

@auto_reset_index
def add_user_statistics(df):
   df_statistics = get_user_statistics(df.userId.unique())
   return df.merge(df_statistics, on='userId', how='left')

加完装饰器后，不管是普通列结构的DataFrame，还是已经把userId/sessionId设为索引的DataFrame，都可以直接用pipe链式调用，不需要手动做索引重置：

# 普通列结构的df可正常运行
df_enriched = df.pipe(add_session_statistics).pipe(add_user_statistics)

# 设了多级索引的X也可直接运行，不需要提前reset_index
X_enriched = X.pipe(add_session_statistics).pipe(add_user_statistics)

内容的提问来源于stack exchange，提问作者MYK