Pandas Pipeline

pandaspipeline.py

import pandas as pd
from inspect import isroutine, signature
from functools import wraps

__all__ = ['pipeline']

if '_all' not in globals():  # Don't redefine this on reloads; otherwise we load the new all
    _all = all  # This will get overridden when we define methods from pd.DataFrame

def is_pandas(obj):
    return isinstance(obj, pd.core.base.PandasObject)

class PandasPipeline():
    def __init__(self, func, *args, **kw):
        self.func = func
        self.args = args
        self.kw = kw
        
    def __rrshift__(self, df):
        if is_pandas(df):
            return self.func(df, *self.args, **self.kw)
        return NotImplemented

def pipeline(n_df=1):
    def decorator(func):
        @wraps(func)
        def decorated(*args, **kw):
            if len(args) >= n_df and _all(is_pandas(arg) for arg in args[:n_df]):
                try:
                    signature(func).bind(*args, **kw)
                except TypeError:
                    # Didn't have the right number of arguments, so try a pipeline
                    pass
                else:
                    return func(*args, **kw)
            return PandasPipeline(func, *args, **kw)
        return decorated
    
    if isroutine(n_df):
        func = n_df
        n_df = 1    # Note that this changes n_df in decorated.
        return decorator(func)
    
    return decorator

def _load_methods():
    def method_pipeline(func):
        @wraps(func)
        def decorated(*args, **kw):
            return PandasPipeline(func, *args, **kw)
        return decorated

    global_dict = globals()

    for method in dir(pd.DataFrame):
        if 'a' <= method[0] <= 'z':
            global_dict[method] = method_pipeline(getattr(pd.DataFrame, method))
            __all__.append(method)

_load_methods()

Presentation.ipynb