pandas magic.py

import inspect
import types

    
def injected(df, thunk):
    """Evaluate a thunk in the context of DataFrame
    
    >>> df = pd.DataFrame({'x': [0, 1, 2]}, index=['a', 'b', 'c'])
    >>> injected(df, lambda: x ** 2)
    a    0
    b    1
    c    4
    Name: x, dtype: int64
    """
    new_globals = thunk.func_globals.copy()
    new_globals.update(df)
    new_thunk = types.FunctionType(thunk.func_code, new_globals, thunk.func_name,
                                   thunk.func_defaults, thunk.func_closure)
    return new_thunk()


def is_thunk(f):
    return callable(f) and not inspect.getargspec(f).args


class MagicFrame(object):
    def __init__(self, df):
        self.__df = df

    def __eval(self, expr):
        if is_thunk(expr):
            expr = injected(self, expr)
        return expr
    
    def __ensure_magical(self, result):
        if isinstance(result, pd.DataFrame):
            result = type(self)(result)
        return result
    
    def __wrapped(method_name):
        def f(self, *args, **kwargs):
            if args:
                args = (self.__eval(args[0]),) + args[1:]
            method = getattr(self.__df, method_name)
            result = method(*args, **kwargs)
            return self.__ensure_magical(result)
        return f

    __getitem__ = __wrapped('__getitem__')
    query = __wrapped('query')
    # TODO: propagate magic through groupby operations
    groupby = __wrapped('groupby')
    describe = __wrapped('describe')
    
    def __getattr__(self, name):
        return getattr(self.__df, name)
    
    def __dir__(self):
        return dir(self.__df)


# example
import seaborn as sns
import pandas as pd

df = MagicFrame(sns.load_dataset("iris"))
df[lambda: sepal_length > 3].groupby(lambda: pd.cut(sepal_width, 5)).mean()