How to use the `.loc` method from pandas on a custom class object?

Viewed 948

I've been going through the source code of pandas https://github.com/pandas-dev/pandas/blob/master/pandas/core/generic.py and I can't figure out where they actually implement the .loc slicing method. I'm working on a wrapper that takes in a bunch of pd.DataFrames. For the sake of this question, let's call it DataFrameCollection. I don't want to inherit all of the methods so I don't want to do class DataFrameCollection(pd.DataFrame): pass.

Does anyone know which code is responsible for the .loc method of a pd.DataFrame object and how this can be used on a custom object?

Essentially I would like to be able to do the following:

dfc_iris =  DataFrameCollection(" a bunch of dataframes")
dfc_iris.loc[idx_obsvs, :]
1 Answers

The loc attribute is one of several indexers, see the pandas.core.indexing module, specifically the get_indexers_list() function:

# the supported indexers
def get_indexers_list():


    return [
        ('ix', _IXIndexer),
        ('iloc', _iLocIndexer),
        ('loc', _LocIndexer),
        ('at', _AtIndexer),
        ('iat', _iAtIndexer),
    ]

Each of those classes is defined in the same module.

That function is used to add attributes to the NDFrame class, which is a base class of pandas.DataFrame. Each of the classes in the get_indexers_list() result is added as a property object.

So to re-use the object type, you could add your properties, using the same code if necessary; add the same class method to your class

@classmethod
def _create_indexer(cls, name, indexer):
    """Create an indexer like _name in the class."""
    if getattr(cls, name, None) is None:
        _indexer = functools.partial(indexer, name)
        setattr(cls, name, property(_indexer, doc=indexer.__doc__))

then add the indexers with

# install the indexes
for _name, _indexer in indexing.get_indexers_list():
    DataFrameCollection._create_indexer(_name, _indexer)

Given a dfcollection instance of your DataFrameCollection class, dfcollection.loc would then result in _LocIndexer('loc', dfcollection) being called and returned.

Do study the remaining code in pandas.core.indexing to see how each indexer then expects to find information on your DataFrameCollection instance; it's the self.obj reference in the indexer methods.

For example, dfcollection.loc[...] is translated to _LocationIndexer.__getitem__(), which delegates to _LocIndexer._is_scalar_access(), _LocIndexer._getitem_scalar(), _NDFrameIndexer._getitem_tuple() and _LocIndexer._getitem_axis(), which together with methods these delegate to, need access to at least the .axes, .ndim ._get_value(), ._get_axis_name(), ._get_axis_number(), ._get_axis(), ._reindex_with_indexers() and ._take() attributes and methods on the dataframe.

Related