1

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

AMC
  • 2,642
  • 7
  • 13
  • 35
semidevil
  • 69
  • 1
  • 11

1 Answers1

2

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]
Partha Mandal
  • 1,391
  • 8
  • 14
  • this works! thank you. What does @property do? and why did we need to add self.dataset in the returns? – semidevil Apr 16 '20 at 22:15
  • 1
    We need `self.dataset` because that actually refers to your `pandas_obj` (and `self` refers to your `my_extension` class). For more on `property`, this is a good link : https://stackoverflow.com/questions/17330160/how-does-the-property-decorator-work – Partha Mandal Apr 16 '20 at 22:38