如何在 Python 中將現有函數包裝在條件邏輯上？

Question

我有一個代碼庫，這種模式很常見：

df # Some pandas dataframe with columns userId, sessionId 

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

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

# etc..

df_enriched = (df
               .pipe(add_session_statistics)
               .pipe(add_user_statistics)
               )

但是，在代碼庫的另一部分中，我將“userId”、“sessionId”作為數據幀的索引。 就像是：

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

這意味着我不能在X上使用add_{somthing}_statistics()函數而不每次都重置索引。

是否有任何裝飾器可以添加到add_{somthing}_statistics()以使他們在嘗試合並不存在的列時遇到KeyError時重置索引？

Answer 1

這似乎有效：

def index_suspension_on_add(add_function):
    
    def _helper(df):
        try:
            return df.pipe(add_function)
        except Exception:
            index_names = df.index.names 

            return (df
                    .reset_index()
                    .pipe(add_function)
                    .set_index(index_names)
            )
            
    return _helper

@index_suspension_on_add
def add_user_statistics(df):
    ...