zab
zab

Reputation: 155

Custom methods for a Pandas DataFrame

I'd like to write custom functions that I can call on my pd.DataFrame df using the df.method() notation.

For instance,

def my_pd_method(df: pd.DataFrame, col: str)->pd.DataFrame:

    '''apply my_function to df[col] of df'''

    df_copy = df.copy(deep = True)
    df_copy[col] = df_copy[col].apply(lambda x: my_function(x), axis = 1)
    
   return df_copy 

After this, I can run the command

PandasObject.my_pd_method = my_pd_method 

to define the my_pd_method as a pd method.

After this, df.my_pd_method(col) will run as expected.

Is there some way to do this in a single function that I can put it in a library, import it and start using it, without having to run PandasObject.my_pd_method = my_pd_method?

Upvotes: 2

Views: 164

Answers (1)

vianmixt
vianmixt

Reputation: 883

Your best shot is to use inheritance i.e to create your own custom class that inherits from pandas DataFrame class.

Example:

class CustomDataFrame(pd.DataFrame):
    def my_method(self, col):
        df_copy = self.copy(deep = True)
        df_copy[col] = df_copy[col].apply(lambda x: my_function(x), axis = 1)

    return df_copy 

Then you will be able to call your method like you wanted:

df.my_method(col)

Upvotes: 2

Related Questions