рефакторинг Python-функции, которая использует кадры данных pandas

Question

Я получил в своем коде 20+ из и, безусловно, будет больше вхождений следующего фрагмента кода, где у меня почти всегда один и тот же шаблон кода, который по сути является (число соответствует комментариям в фрагменте ниже):

1. вычисляет тест, который зависит от нескольких столбцов df, тест может умножать, делить, добавлять любые операции над любыми столбцами и заменять inf на nan
2. , чтобы получить маску для всех non nan
3. создайте действительный df с помощью маски
4. создайте новый столбец, добавив '_mod' к рассматриваемому исходному столбцу и заполнив его чем-то
5. , заполнив остальные значения встолбец «_mod» с другой формулой, которая, как и в шаге 1, может быть любой операцией над любыми столбцами, используя действительный df, ранее созданный в 3

Фрагмент

# 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.Тем не менее, я не вижу, как передать список операций в качестве параметра в рефакторированную функцию, которая использует somehting (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')

Как правило, вам никогда не нужно пропускать series в качестве аргументов функции.