Making IPython Tab-Complete Hierarchical pandas DataFrames

Question

One of the useful features of IPython is its tab completion, among other things eliminating the need to memorize pandas.DataFrame column names.

E.g., suppose we have

df = pd.DataFrame({'bar': [1, 2], 'baz': [3, 4], 'bap': [5, 6]})

giving df as

    bap bar baz
0   5   1   3
1   6   2   4

Then we can type

df.<Tab>

and it will show bap, bar, and baz (among others) as members, as well as attempting to complete them.

Unfortunately, this useful functionality partially disappears with hierarchical dataframes. E.g., if we change things to

df = pd.DataFrame({
    ('foo', 'bar'): [1, 2], 
    ('foo', 'baz'): [3, 4], 
    ('bap', ''): [5, 6]})

giving df as

    bap foo
        bar baz
0   5   1   3
1   6   2   4

then df.<Tab> will not autocomplete bap or foo.

What is the proper way to get this working? I have written a hack which does this (answer below), but am displeased with its reliance on Python monkey-patching. Other answers would be welcome.

Answer 1

Applying the following monkey-patches seems to remove the symptoms of the problem:

_orig_dir = getattr(pd.DataFrame, '__dir__')
def _mutilevel_aware_dir(df):
    multilevels = [tup for tup in df.columns if isinstance(tup, tuple)]
    return _orig_dir(df) + [tup[0] for tup in multilevels]
setattr(pd.DataFrame, '__dir__', _mutilevel_aware_dir)


_orig_get_attr = getattr(pd.DataFrame, '__getattr__')
def _mutilevel_aware_getattr(df, name):
    return _orig_get_attr(df, name)
setattr(pd.DataFrame, '__getattr__', _mutilevel_aware_getattr)

IPython apparently indirectly uses the __dir__ method of a DataFrame for autocompletion. The first function internally scan for tuples, and if they are encountered, returns their first elements as "dummy" members. The second function grabs the __getattr__ for handling what will be returned when such a dummy member is accessed.

Patching these into the class seems to do the job. However, this is a rather violent solution, and it might have side effects of which I'm unaware.