Using iloc and loc from pandas in own class-CodePudding

I want to make an own class ExampleClass, that stores data in a pandas DataFrame. I do this to create my own methods to manipulate the data.

My question is: Is it possible, to use an object of ExampleClass like:

import pandas as pd

class ExampleClass:
    def __init__(self, data: pd.DataFrame):
        self.data = data
        self.other_value = 42

obj = ExampleClass(pd.DataFrame([1,2,3]))
part_of_obj = obj.iloc[:2]

where type(part_of_obj) == ExampleClass and part_of_obj.data == obj.data.iloc[:2]?

So I want a copy of obj but with obj.data.iloc[:2] .

I don't want to use __getitem__ because then I'm bound to iloc OR loc and cannot use both. And I would prefer not to make an own method called iloc(), because I couldn't access it easy with [].

Maybe I could make it with inheritance, but I've already written my class and it uses a DataFrame to manipulate the data, what would make it very difficult to change everything written.

CodePudding user response：

You can expose iloc via a property then you can also add loc the same way to subset the dataframe.

@property
def iloc(self):
  return self.data.iloc

To return an object of the same type maybe a helper class can be used that wraps the call to __getitem__:

class ExampleClass:
  class Subsetter:
    def __init__(self, locator):
      self.locator = locator

    def __getitem__(self, *vargs, **kwargs):
      return ExampleClass(self.locator.__getitem__(*vargs, **kwargs))

  def __init__(self, data: pd.DataFrame):
      self.data = data
      self.other_value = 42

  @property
  def iloc(self):
    return self.Subsetter(self.data.iloc)

obj = ExampleClass(pd.DataFrame([1,2,3]))
part_of_obj = obj.iloc[:2]