CODEWITHSUNDEEP
pythonpandas
semidevil
semidevil

Reputation: 69

having trouble extending pandas

I'm trying the "Extend pandas" functionality but can't seem to get it to work. I'm trying to follow the example provided by pandas.

Link

Heres my code so far with fake data

import pandas as pd
data =pd.DataFrame({'ID1':[1,2], 'ID2':[3,4], 'value':[5,6]})
@pd.api.extensions.register_dataframe_accessor("ext")
class my_extension:
    def __init__(self, pandas_obj):
        self.dataset = pandas_obj
    @property
    def retrieve_ID(self, ID):
        return self[self['ID1']==ID]

    def method2(self, ID):
        return self[self['ID2']==ID]

Basically, the end result that I want is for me to type "data.ext." and hit tab and see retrieve_ID and method2 as an option for my functions

Upvotes: 1

Views: 468

Answers (1)

Partha Mandal
Partha Mandal

Reputation: 1441

Pretty sure you are missing the self.dataset part (and don't think you need the @property there). Btw this should work:

import pandas as pd
data =pd.DataFrame({'ID1':[1,2], 'ID2':[3,4], 'value':[5,6]})
@pd.api.extensions.register_dataframe_accessor("ext")
class my_extension:
    def __init__(self, pandas_obj):
        self.dataset = pandas_obj
    #@property
    def retrieve_ID(self, ID):
        return self.dataset[self.dataset['ID1']==ID]

    def method2(self, ID):
        return self.dataset[self.dataset['ID2']==ID]

Upvotes: 2

Related Questions