columns_in_formula function that parses all of the column names from a patsy formula.

patsycolumns.py

from StringIO import StringIO
from tokenize import generate_tokens, NAME

import patsy


def _tokens_from_patsy(node):
    """
    Yields all the individual tokens from within a patsy formula
    as parsed by patsy.parse_formula.parse_formula.

    Parameters
    ----------
    node : patsy.parse_formula.ParseNode

    """
    for n in node.args:
        for t in _tokens_from_patsy(n):
            yield t

    if node.token:
        yield node.token


def columns_in_formula(formula):
    """
    Returns the names of all the columns used in a patsy formula.

    Parameters
    ----------
    formula : str
        A patsy formula.

    Returns
    -------
    columns : list of str

    """
    columns = []

    tokens = map(
        lambda x: x.extra,
        filter(
            lambda x: x.extra is not None,
            _tokens_from_patsy(patsy.parse_formula.parse_formula(formula))))

    for tok in tokens:
        # if there are parentheses in the expression we
        # want to drop them and everything outside
        # and start again from the top
        if '(' in tok:
            start = tok.find('(') + 1
            fin = tok.rfind(')')
            columns.extend(columns_in_formula(tok[start:fin]))
        else:
            for toknum, tokval, _, _, _ in generate_tokens(
                    StringIO(tok).readline):
                if toknum == NAME:
                    columns.append(tokval)

    return list(set(columns))