generate partial polynomial features

def partial_poly(data=[], cols=[], deg=2, interaction=False):
    # data=[train_data,test_data, ...]
    # cols=['selected','cols','to','transform']
    poly = PolynomialFeatures(
        degree=deg, interaction_only=interaction, include_bias=False)
    out_df = []
    for i in range(len(data)):
        poly_ = pd.DataFrame(poly.fit_transform(data[i][cols]),
                             columns=poly.get_feature_names_out(),
                             index=data[i].index,
                             ).drop(columns=cols)
        out_df.append(pd.concat([data[i], poly_], axis=1))
    return out_df

poly_transformed = partial_poly(data=[X_train, X_test], cols=['A','B','C'])
poly_transformed[0] # transformed X_train
poly_transformed[1] # transformed X_test

입력되는 데이터 셋 모두에 cols=[] 컬럼이 유효해야 함.