pipe decorator example

pdf_pipe_decorator.py

def cell_wrapper(df, func, field, drop=True, new_name=None):
    """
    decorator function for pandas pipe api
    takes func which applies function to one value in field
    returns modified dataframe
    df (pandas dataframe): the dataframe to apply transformation on
    func (function): function to apply to each value of field
    field (str): name of column in df
    drop (boolean): whether to drop 'field' after transformation
    new_name (str): whether to rename transformed 'field' column to new_name
    """
    lst = list(df[field].apply(func))
    if isinstance(lst[0], dict):
        adf = pd.DataFrame(lst, index=df.index)
    else:
        adf = pd.DataFrame(lst, index=df.index, columns=[field])
    if new_name is not None:
        adf = adf.rename(columns={field: new_name})
    if drop:
        return pd.concat([df.drop(field, axis=1), adf], axis=1)
    return pd.concat([df, adf], axis=1)

def an_apply_func(value):
    # do something complicated to value
    new_value = value
    return new_value

from functools import partial

pipe_func = partial(cell_wrapper, func=an_apply_func, field='column_name')


def df_wrapper(df, func):
"""
decorator function for pandas dataframe pipe api
takes function that transforms df and then concats by column and returns both original and transformed
df (pandas dataframe): dataframe of interest
func (function): function to apply to dataframe to produce new dataframe to concat horizontally
"""
new_df = func(df)
return pd.concat([df, new_df], axis=1)