python - 重构使用 pandas 数据帧的 python 函数

Question

我的代码中有超过 20 个，并且肯定会出现更多以下代码片段，其中我几乎总是具有相同的代码模式，其本质是(数字对应于下面代码片段中的注释):

1. 计算一个依赖于 df 的几列的测试，该测试可以对任何列进行乘法、除法、加法、任何操作，并将 inf 替换为 nan
2. 戴上所有非楠的面具
3. 使用掩码创建有效的 df
4. 创建一个新列，将“_mod”添加到所考虑的原始列中并用某些内容填充它
5. 使用第 1 步中的另一个公式填充“_mod”列上的其余值，可以是使用之前在第 3 步中创建的有效 df 对任何列进行的任何操作

代码片段

# col1
logger.info('col1')
# 1
col1_test = (df["colX"] / df["colZ"] / df["colY"] / df["colX"]).replace([np.inf, -np.inf], np.nan)
# 2
col1_mask = (~pd.isna(col1_test))
# 3
col1_valid = df[col1_mask]
# 4
df['col1_mod'] = np.nan
# 5
df.loc[col1_mask, 'col1_mod'] = (col1_valid["colX"] - col1_valid["colZ"]) / col1_valid[
    "colY"]

# col2
logger.info('col2')
col2_test = (df["colA"] / df["colY"] / df["colA"]).replace(
    [np.inf, -np.inf], np.nan)
col2_mask = (~pd.isna(col2_test))
col2_valid = df[col2_mask]
df['col2_mod'] = 0.0
df.loc[col2_mask, 'col2_mod'] = col2_valid["colA"] / col2_valid["colY"]

到目前为止我所写的重构内容如下。但我想它可以更进一步，我特别被注释的 (# df.loc[mask, f'{oldcol}_mod'] = ...) 阻止，这可能如果函数返回 df 本身，则解决整个问题。然而，我不知道如何将操作列表作为参数传递给使用重构函数本身创建的某些内容(valid)的重构函数。

def refactored(df, oldcol, dftest, replace):
    logger.info(oldcol)
    test = dftest.replace([np.inf, -np.inf], np.nan)
    mask = (~pd.isna(test))
    valid = df[mask]
    df[f'{oldcol}_mod'] = replace
    # df.loc[mask, f'{oldcol}_mod'] = ...
    return valid, mask


col1_valid, col1_mask = refactored(df, 'col1', df["colX"] / df["colZ"] / df["colY"] / df["colX"], np.nan)
df.loc[col1_mask, 'col1_mod'] = (col1_valid["colX"] - col1_valid["colZ"]) / col1_valid["colY"]
col2_valid, col2_mask = refactored(df, 'col2',df["colA"] / df["colY"] / df["colA"] , 0.0)
df.loc[col2_mask, 'col2_mod'] = col2_valid["colA"] / col2_valid["colY"]

Answer 1

考虑使用pd.DataFrame.eval与 pd.DataFrame.pipe :

def refactored(df, oldcol, dftest, replace, mod_col, series_col):
    # ...some logic...
    test = df.eval(dftest).replace([np.inf, -np.inf], np.nan)
    # ... some more logic...
    df.loc[mask, mod_col] = df.eval(series_col)
    return df

df = df.pipe(refactored, 'col1', 'colX / colZ / colY / colX', np.nan,
             'col1_mod', '(colX - colZ) / colY')\
       .pipe(refactored, 'col2', 'colA / colY / colA', 0.0,
             'col2_mod', 'colA / colY')

一般来说，您永远不需要将系列作为函数参数传递。