I want to make a function that does an ETL to a list/Series and returns a Series with additional attributes and methods specific to that function. I can achieve this by creating a class to extend the Series and it works but then when I try to reassign the output from the function with the new class the updated class attributes and methods are stripped away. How can I extend the Series to have custom attributes and methods that are not stripped away when reassigning back to the dataframe?
Custom function that does ETL, returns a Series with an extended class
import pandas as pd
def normalize_x(x: list, new_attribute: None):
normalized = pd.Series(['normalized_'+ i if i != 4 else None for i in x])
return NormalizeX(normalized = normalized, original = x, new_attribute = new_attribute)
class NormalizeX(pd.Series):
def __init__(self, normalized, original, new_attribute, *args, **kwargs,):
super().__init__(data = normalized, *args, **kwargs)
self.original = original
self.normalized = normalized
self.new_attribute = new_attribute
def conversion_errors(self):
return [o != n for o, n in zip(pd.isnull(self.original), pd.isnull(self.normalized))]
df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": ['dog', 'cat', 4]})
Assign to a new object (new attributes and methods work)
out = normalize_x(df.C, new_attribute = 'CoolAttribute')
out
## 0 normalized_dog
## 1 normalized_cat
## 2 None
## dtype: object
## Can still use Series methods
out.to_list()
## ['normalized_dog', 'normalized_cat', None]
## Can use the new methods and access attributes
out.conversion_errors()
## [False, False, True]
out.original
##0 dog
##1 cat
##2 4
##Name: C, dtype: object
Assign to a Pandas DataFrame (new attributes and methods break)
df['new'] = normalize_x(df.C, new_attribute = 'CoolAttribute')
df['new']
## 0 normalized_dog
## 1 normalized_cat
## 2 None
## dtype: object
## Can't use the new methods or access attributes
df['new'].conversion_errors()
## AttributeError: 'Series' object has no attribute 'conversion_errors'
df['new'].original
## AttributeError: 'Series' object has no attribute 'original'