DataFrame pipeable function decorator

pipeable_function_decorator.py

from functools import wraps

def pipeable_func(func):
    """allows any function to be used in a 
    ``.pipe`` context on a dataframe, assigning output from this function as applied 
    to a dataframe column.
    Usage:
    
        @pipeable_func
        def is_odd(n):
            return n % 2


        (pd.DataFrame(np.random.randint(1, high=100000, size=5), columns=["a_number"])
         .pipe(is_odd, "id")
        )

    """
    @wraps(func)
    def wrapper(df, apply_col_name, assign_name=None):
        name = func.__name__ if assign_name is None else assign_name
        return df.assign(**{name: df[apply_col_name].apply(func)})
    return wrapper
        
    
@pipeable_func
def is_odd(n):
    return n % 2


def is_odd_long(n):
    return n % 2


## Example of usage - see the difference between the two methods of assigning columns in a method-chaning manner

(pd.DataFrame(np.random.randint(1, high=100000, size=5), columns=["id"])
 .pipe(is_odd, "id")
 .assign(is_odd_long=lambda df: df["id"].apply(is_odd_long))
)