扩展 pandas 时遇到问题

Question

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

Answer 1

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]