wrap pandas daraframe within class

Question

I'd like to give pandas dataframe a custom and extend it with custom methods, but still being able to use the common pandas syntax on it.

I implemented this:

import pandas as pd

class CustomDF():
    def __init__(self, df: pd.DataFrame):
        self.df = df

    def __getattr__(self, name):
        return getattr(self.df, name)

    def foo(self):
        print('foo')
        return

the problem with the above code is that I'd like to make the following lines work:

a = CustomDF()
a.iloc[0,1]
a.foo()

but if I try to access a column of the dataframe by doing

print(a['column_name'])

I get the error "TypeError: 'CustomDF' object is not subscriptable"

Any idea on how not requiring to access .df first to get the subscription working for the super class? Thanks!

why don't you use inheritance? `class CustomDF(pd.DataFrame): `? — ansev, Jan 20 '20 at 15:21
unfortunately I get the error "RecursionError: maximum recursion depth exceeded while calling a Python object" — user1403546, Jan 20 '20 at 15:24
what you really want is an extension of the class, that is, you want the new object to have the methods of `pd.DataFrame`, so the solution must be inheritance — ansev, Jan 20 '20 at 15:26
sorry, my question was wrong. The problem is only that the class is not subscriptable with square brackets... see updated question — user1403546, Jan 20 '20 at 15:42
There's a section on extending Pandas classes in the documentation which shows various approaches to achieve what you want: https://pandas.pydata.org/pandas-docs/stable/development/extending.html — Pete C, Jun 25 '21 at 09:05

score 2 · Answer 1 · answered Oct 29 '22 at 19:04

2

Try inheriting from the pandas.DataFrame class. Like so:

from pandas import DataFrame

class CustomDF(DataFrame):
    def foo(self):
        print('foo')
        return
    
a = CustomDF([0,1])
a.foo()

answered Oct 29 '22 at 19:04

abgo2020

175
1
1
10

wrap pandas daraframe within class

1 Answers1